code: def LSTM(
      n_units, mode="train", return_state=False, initial_state=False
  ):
      """LSTM running on axis 1.
  
      Args:
        n_units: `n_units` for the `LSTMCell`.
        mode: if 'predict' then we save the previous state for one-by-one inference.
        return_state: Boolean. Whether to return the latest status in addition to
          the output. Default: False.
        initial_state: Boolean. If the state RNN (c, h) is to be obtained from the
          stack. Default: False.
  
      Returns:
        A LSTM layer.
      """
  
      if not initial_state:
          zero_state = MakeZeroState(
              depth_multiplier=2
          )  # pylint: disable=no-value-for-parameter
          if return_state:
              return cb.Serial(
                  cb.Branch([], zero_state),
                  cb.Scan(LSTMCell(n_units=n_units), axis=1, mode=mode),
                  name=f"LSTM_{n_units}",
                  sublayers_to_print=[],
              )
          else:
              return cb.Serial(
                  cb.Branch(
                      [], zero_state
                  ),  # fill state RNN with zero.
                  cb.Scan(LSTMCell(n_units=n_units), axis=1, mode=mode),
                  cb.Select([0], n_in=2),  # Drop RNN state.
                  # Set the name to LSTM and don't print sublayers.
                  name=f"LSTM_{n_units}",
                  sublayers_to_print=[],
              )
      else:
          if return_state:
              return cb.Serial(
                  cb.Scan(LSTMCell(n_units=n_units), axis=1, mode=mode),
                  name=f"LSTM_{n_units}",
                  sublayers_to_print=[],
              )
          else:
              return cb.Serial(
                  cb.Scan(LSTMCell(n_units=n_units), axis=1, mode=mode),
                  cb.Select([0], n_in=2),  # Drop RNN state.
                  name=f"LSTM_{n_units}",
                  sublayers_to_print=[],
              )
  
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 a global function with example usage: LSTM(n_units, mode, return_state, initial_state) and returns: cb, cb, cb, cb