code: def test_lstm_ref(
      self,
      batch_size: int,
      seq_len: int,
      input_size: int,
      hidden_size: int,
      num_layers: int,
      bidirectional: bool,
  ):
      num_directions = 2 if bidirectional else 1
  
      seq_lengths = jax.random.randint(
          jax.random.PRNGKey(0),
          (batch_size,),
          0,
          seq_len,
          dtype=jnp.int32,
      )
  
      root_key = jax.random.PRNGKey(1)
      k1, k2, k3, k4 = jax.random.split(root_key, 4)
      x = jax.random.normal(
          k1, (batch_size, seq_len, input_size), dtype=jnp.float32
      )
      h_0 = jax.random.normal(
          k2,
          (num_directions * num_layers, batch_size, hidden_size),
          dtype=jnp.float32,
      )
      c_0 = jax.random.normal(
          k3,
          (num_directions * num_layers, batch_size, hidden_size),
          dtype=jnp.float32,
      )
      weights = rnn.init_lstm_weight(
          k4, input_size, hidden_size, num_layers, bidirectional
      )
  
      @partial(jax.value_and_grad, has_aux=True)
      def f(weights, x, h_0, c_0):
          W_ih, W_hh, b_ih, b_hh = rnn.unpack_lstm_weights(
              weights,
              input_size,
              hidden_size,
              num_layers,
              bidirectional,
          )
          y_ref, h_n_ref, c_n_ref = rnn.lstm_ref(
              x,
              h_0,
              c_0,
              W_ih,
              W_hh,
              b_ih,
              b_hh,
              seq_lengths=seq_lengths,
              input_size=input_size,
              hidden_size=hidden_size,
              num_layers=num_layers,
              dropout=False,
              bidirectional=bidirectional,
          )
  
          loss = jnp.sum(y_ref)
  
          return loss, (y_ref, h_n_ref, c_n_ref)
  
      (loss_ref, (y_ref, h_n_ref, c_n_ref)), grad_ref = f(
          weights, x, h_0, c_0
      )
  
      self.assertFalse(np.isnan(loss_ref))
      self.assertFalse(np.isnan(grad_ref).any())
  
      self.assertEqual(
          y_ref.shape,
          (batch_size, seq_len, num_directions * hidden_size),
      )
  
      for i in range(batch_size):
          y_padded = y_ref[i, seq_lengths[i] :]
          np.testing.assert_allclose(y_padded, jnp.zeros_like(y_padded))
  
file_path: jax/tests/experimental_rnn_test.py
imports: from functools import partial
  from absl.testing import absltest
  import numpy
  import jax
  import jax.numpy
  from jax._src import config
  from jax._src import test_util
  from jax.experimental import rnn
label_desc: Use image node_id 2 for calling the RnnTest obj's underlying member method code with example usage: obj.test_lstm_ref(batch_size, seq_len, input_size, hidden_size, num_layers, bidirectional) and returns: loss,  y_ref,  h_n_ref,  c_n_ref