code: def test_output_shapes(
      self, time_step, input_dim, output_dim, hidden_init
  ):
      obs_inputs = np.full((self.batch_size, time_step, input_dim), 1.0)
      obs_input = np.full((self.batch_size, input_dim), 1.0)
  
      input_var = tf.compat.v1.placeholder(
          tf.float32, shape=(None, None, input_dim), name="input"
      )
      step_input_var = tf.compat.v1.placeholder(
          tf.float32, shape=(None, input_dim), name="step_input"
      )
      output_nonlinearity = tf.keras.layers.Dense(
          units=output_dim,
          activation=None,
          kernel_initializer=tf.constant_initializer(1),
      )
      with tf.compat.v1.variable_scope("GRU"):
          self.gru = gru(
              all_input_var=input_var,
              name="gru",
              gru_cell=self.gru_cell,
              step_input_var=step_input_var,
              step_hidden_var=self.step_hidden_var,
              hidden_state_init=tf.constant_initializer(hidden_init),
              output_nonlinearity_layer=output_nonlinearity,
          )
  
      self.sess.run(tf.compat.v1.global_variables_initializer())
  
      # Compute output by doing t step() on the gru cell
      outputs_t, output_t, h_t, hidden_init = self.gru
      hidden = np.full(
          (self.batch_size, self.hidden_dim), hidden_init.eval()
      )
  
      for _ in range(time_step):
          output, hidden = self.sess.run(
              [output_t, h_t],
              feed_dict={
                  step_input_var: obs_input,
                  self.step_hidden_var: hidden,
              },
          )  # noqa: E126
          assert output.shape == (self.batch_size, output_dim)
          assert hidden.shape == (self.batch_size, self.hidden_dim)
  
      full_output = self.sess.run(
          outputs_t, feed_dict={input_var: obs_inputs}
      )
  
      assert full_output.shape == (
          self.batch_size,
          time_step,
          output_dim,
      )
  
file_path: garage/tests/garage/tf/models/test_gru.py
imports: import numpy
  import pytest
  import tensorflow
  from garage.tf.models.gru import gru
  from tests.fixtures import TfGraphTestCase
  from tests.helpers import recurrent_step_gru
label_desc: Use image node_id 2 for calling the TestGRU obj's underlying member method code with example usage: obj.test_output_shapes(time_step, input_dim, output_dim, hidden_init) without return types