How to optimize LSTM with jax.lax.scan?

[Sun-Xiaohui]

Hi, I created BiLSTM by applying LSTMCell in a "for loop", here LSTMCell is implemented based on JAX. When I apply LSTM on a sequence, I need process each element in the sequence in a "for loop", just like this:

seq_outputs = []
for t in range(seq_len):
        out_forward = modulelist[0]([input_seq[:, t, :], (h[:num_layers, :, :], c[:num_layers, :, :])])
        out_back = modulelist[1]([input_seq[:, seq_len - t - 1, :], (h[num_layers:, :, :], c[num_layers:, :, :])])
        out = jnp.concatenate((out_forward, out_back), axis=-1)
        seq_outputs.append(out)

modulelist[i] is a Sequential of multiple LSTMCell. Jit compilation takes a long time due to the seq_len is large. If I rewrite the code by using jax.lax.scan(), like this:

def help(seq_len, modulelist, h, c, num_layers, input_seq):
  seq_outputs = []
  carry = (seq_len, modulelist, h, c, num_layers, input_seq, seq_outputs)
  
  def loop(carry, t):
    seq_len, modulelist, h, c, num_layers, input_seq, seq_outputs = carry
    hidden_forward, out_forward = modulelist[0]([input_seq[:, t, :], (h[:num_layers, :, :], c[:num_layers, :, :])])
    hidden_back, out_back = modulelist[1]([input_seq[:, seq_len - t - 1, :], (h[num_layers:, :, :], c[num_layers:, :, :])])
    out = jnp.concatenate((out_forward, out_back), axis=-1)
    seq_outputs.append(out)
    return (seq_len, modulelist, h, c, num_layers, input_seq, seq_outputs), (seq_outputs, hidden_forward, hidden_back)

  return jax.lax.scan(loop, carry, jnp.arange(seq_len))[1]

The error occurs: Value <__main__.Sequential object at 0x7f1c3c342340> with type <class '__main__.Sequential'> is not a valid JAX type
I know it's because in my own frameworks, Module is not a PyTree type, leading to Sequential and LSTMCell are not PyTree type(Sequential and LSTMCell inherit from Module ). I wonder how can I speeded the code without modify my Module？ Appreciate！