class_docstr: LSTM Cell.
  
  For a nice overview of the motivation and (i, o, f) gates, see this tutorial:
  https://colah.github.io/posts/2015-08-Understanding-LSTMs/
  
  See this paper for a description and detailed study of all gate types:
  https://arxiv.org/pdf/1503.04069.pdf
class_docstr_tok: ["LSTM","Cell",".","For","a","nice","overview","of","the","motivation","and","(","i",",","o",",","f",")","gates",",","see","this","tutorial",":","https",":","\/\/colah.github.io\/posts\/2015-08-Understanding-LSTMs\/","See","this","paper","for","a","description","and","detailed","study","of","all","gate","types",":","https",":","\/\/arxiv.org\/pdf\/1503.04069.pdf"]
code: def forward(self, inputs):
      x, lstm_state = inputs
  
      # LSTM state consists of c and h.
      c, h = jnp.split(lstm_state, 2, axis=-1)
  
      # Dense layer on the concatenation of x and h.
      w, b = self.weights
      y = jnp.dot(jnp.concatenate([x, h], axis=-1), w) + b
  
      # i = input_gate, j = new_input, f = forget_gate, o = output_gate
      i, j, f, o = jnp.split(y, 4, axis=-1)
  
      new_c = c * fastmath.sigmoid(f) + fastmath.sigmoid(i) * jnp.tanh(
          j
      )
      new_h = jnp.tanh(new_c) * fastmath.sigmoid(o)
      return new_h, jnp.concatenate([new_c, new_h], axis=-1)
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import activation_fns
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
label_desc: Use image node_id 2 for calling the LSTMCell obj's underlying member method code with example usage: obj.forward(inputs) and returns: new_h,  jnp