code: def testWrapperCall(self):
      wrapper = rnn_wrapper.RNNWrapper(
          tf.keras.layers.LSTM(
              3, return_state=True, return_sequences=True
          )
      )
  
      batch_size = 2
      input_depth = 5
      inputs = np.random.rand(batch_size, input_depth).astype(
          np.float32
      )
  
      # Make sure wrapper call works when no time dimension is passed in.
      outputs, next_state = wrapper(inputs)
  
      inputs_time_dim = tf.expand_dims(inputs, axis=1)
      outputs_time_dim, next_state_time_dim = wrapper(inputs_time_dim)
      outputs_time_dim = tf.squeeze(outputs_time_dim, axis=1)
  
      outputs_manual_state, next_state_manual_state = wrapper(
          inputs, wrapper.get_initial_state(inputs)
      )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      for out_variant in (
          outputs,
          outputs_time_dim,
          outputs_manual_state,
      ):
          self.assertEqual(out_variant.shape, (batch_size, 3))
      for state_variant in (
          next_state,
          next_state_time_dim,
          next_state_manual_state,
      ):
          self.assertLen(state_variant, 2)
          self.assertEqual(state_variant[0].shape, (batch_size, 3))
          self.assertEqual(state_variant[1].shape, (batch_size, 3))
  
      self.assertAllClose(outputs, outputs_time_dim)
      self.assertAllClose(outputs, outputs_manual_state)
      self.assertAllClose(next_state, next_state_time_dim)
      self.assertAllClose(next_state, next_state_manual_state)
  
file_path: agents/tf_agents/keras_layers/rnn_wrapper_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow.compat.v2
  from tf_agents.keras_layers import rnn_wrapper
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the RNNWrapperTest obj's underlying member method code with example usage: obj.testWrapperCall() without return types