class_docstr: A simple bi-directional LSTM.
code: def __call__(self, embedded_inputs, lengths):
      # Forward LSTM.
      initial_state = self.forward_lstm.initialize_carry(
          embedded_inputs[:, 0].shape
      )
      _, forward_outputs = self.forward_lstm(
          initial_state, embedded_inputs
      )
  
      # Backward LSTM.
      reversed_inputs = flip_sequences(embedded_inputs, lengths)
      initial_state = self.backward_lstm.initialize_carry(
          reversed_inputs[:, 0].shape
      )
      _, backward_outputs = self.backward_lstm(
          initial_state, reversed_inputs
      )
      backward_outputs = flip_sequences(backward_outputs, lengths)
  
      # Concatenate the forward and backward representations.
      outputs = jnp.concatenate([forward_outputs, backward_outputs], -1)
      return outputs
  
imports: import functools
  from typing import Any, Callable, Optional
  from flax import linen
  import jax
  from jax import numpy
label_desc: Use image node_id 2 for calling the SimpleBiLSTM obj's underlying member method code with example usage: obj.__call__(embedded_inputs, lengths) and returns: outputs