code: def test_lstm(
      self,
      batch_size: int,
      seq_len: int,
      input_size: int,
      hidden_size: int,
      num_layers: int,
      bidirectional: bool,
  ):
      # TODO(phawkins): Partially disable this on cudnn version per b/281071013
      if (
          batch_size == 1
          and seq_len == 4
          and input_size == 1
          and hidden_size == 6
          and num_layers == 4
          and bidirectional == False
      ):
          self.skipTest("Test requires cudnn >= 8.8")
  
      num_directions = 2 if bidirectional else 1
      seq_length_key, root_key = jax.random.split(jax.random.PRNGKey(0))
  
      seq_lengths = jax.random.randint(
          seq_length_key, (batch_size,), 1, seq_len, dtype=jnp.int32
      )
  
      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
      )
  
      def f(weights, x, h_0, c_0):
          y, h, c = rnn.lstm(
              x,
              h_0,
              c_0,
              weights,
              seq_lengths=seq_lengths,
              input_size=input_size,
              hidden_size=hidden_size,
              num_layers=num_layers,
              dropout=False,
              bidirectional=bidirectional,
          )
          seq_length_mask = (
              jnp.tile(
                  jnp.arange(seq_len, dtype=jnp.int32)[None],
                  [batch_size, 1],
              )
              < seq_lengths[:, None]
          )
          loss = jnp.sum(jnp.where(seq_length_mask[..., None], y, 0.0))
          return loss, (y, h, c)
  
      jtu.check_grads(
          f,
          (weights, x, h_0, c_0),
          modes=["rev"],
          order=1,
          atol=5e-3,
          rtol=5e-3,
      )
  
      (loss, (y, h_n, c_n)), weights_grad = jax.value_and_grad(
          f, has_aux=True
      )(weights, x, h_0, c_0)
  
      def g(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,
          )
          seq_length_mask = (
              jnp.tile(
                  jnp.arange(seq_len, dtype=jnp.int32)[None],
                  [batch_size, 1],
              )
              < seq_lengths[:, None]
          )
          loss = jnp.sum(
              jnp.where(seq_length_mask[..., None], y_ref, 0.0)
          )
          return loss, (y_ref, h_n_ref, c_n_ref)
  
      (
          loss_ref,
          (y_ref, h_n_ref, c_n_ref),
      ), weights_grad_ref = jax.value_and_grad(g, has_aux=True)(
          weights, x, h_0, c_0
      )
  
      self.assertAllClose(
          weights_grad_ref, weights_grad, rtol=1e-5, atol=1e-5
      )
      np.testing.assert_allclose(loss_ref, loss, rtol=1e-05, atol=1e-5)
      np.testing.assert_allclose(y_ref, y, rtol=1e-05, atol=1e-5)
      np.testing.assert_allclose(h_n_ref, h_n, rtol=1e-05, atol=1e-5)
      np.testing.assert_allclose(c_n_ref, c_n, rtol=1e-05, atol=1e-5)
  
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 1 for calling the RnnTest obj's underlying member method code with example usage: obj.test_lstm(batch_size, seq_len, input_size, hidden_size, num_layers, bidirectional) and returns: loss,  y,  h,  c, loss,  y_ref,  h_n_ref,  c_n_ref