class_docstr: EncoderLSTM Module wrapped in a lifted scan transform.
class_docstr_tok: ["EncoderLSTM","Module","wrapped","in","a","lifted","scan","transform","."]
docstr: Applies the module.
docstr_tok: ["Applies","the","module","."]
code: def __call__(
      self, carry: tuple[Array, Array], x: Array
  ) -> tuple[tuple[Array, Array], Array]:
      """Applies the module."""
      lstm_state, is_eos = carry
      new_lstm_state, y = nn.LSTMCell(features=lstm_state[0].shape[-1])(
          lstm_state, x
      )
  
      # Pass forward the previous state if EOS has already been reached.
      def select_carried_state(new_state, old_state):
          return jnp.where(is_eos[:, np.newaxis], old_state, new_state)
  
      # LSTM state is a tuple (c, h).
      carried_lstm_state = tuple(
          select_carried_state(*s)
          for s in zip(new_lstm_state, lstm_state)
      )
      # Update `is_eos`.
      is_eos = jnp.logical_or(is_eos, x[:, self.eos_id])
      return (carried_lstm_state, is_eos), y
  
imports: import functools
  from typing import Any
  from flax import linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the EncoderLSTM obj's underlying member method code with example usage: obj.__call__(carry, x) and returns: carried_lstm_state,  is_eos,  y, jnp