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=[])
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.
  
trax/layers/rnn.py