code: def test_rnn_basic_forward(self):
      batch_size = 10
      seq_len = 40
      channels_in = 5
      channels_out = 15
  
      rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)
  
      xs = jnp.ones((batch_size, seq_len, channels_in))
      variables = rnn.init(jax.random.key(0), xs)
      ys: jnp.ndarray
      carry, ys = rnn.apply(variables, xs)
  
      self.assertEqual(ys.shape, (batch_size, seq_len, channels_out))
  
      for layer_params in variables["params"]["cell"].values():
          if "bias" in layer_params:
              self.assertEqual(
                  layer_params["bias"].shape, (channels_out,)
              )
          self.assertIn(
              layer_params["kernel"].shape[0],
              [channels_in, channels_out],
          )
          self.assertEqual(
              layer_params["kernel"].shape[1], channels_out
          )
  
imports: import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
  from flax import linen
  from flax.linen.recurrent import flip_sequences
label_desc: Use image node_id 1 for calling the RNNTest obj's underlying member method code with example usage: obj.test_rnn_basic_forward() without return types