code: def test_output_shapes(
      self, time_step, input_dim, output_dim, hidden_init, cell_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="input"
      )
      _output_nonlinearity = tf.keras.layers.Dense(
          units=output_dim,
          activation=None,
          kernel_initializer=tf.constant_initializer(1),
      )
      with tf.compat.v1.variable_scope("LSTM"):
          self.lstm = lstm(
              all_input_var=_input_var,
              name="lstm",
              lstm_cell=self.lstm_cell,
              step_input_var=_step_input_var,
              step_hidden_var=self._step_hidden_var,
              step_cell_var=self._step_cell_var,
              hidden_state_init=tf.constant_initializer(hidden_init),
              cell_state_init=tf.constant_initializer(cell_init),
              output_nonlinearity_layer=_output_nonlinearity,
          )
  
      self.sess.run(tf.compat.v1.global_variables_initializer())
  
      # Compute output by doing t step() on the lstm cell
      outputs_t, output_t, h_t, c_t, hidden_init, cell_init = self.lstm
      hidden = np.full(
          (self.batch_size, self.hidden_dim), hidden_init.eval()
      )
      cell = np.full(
          (self.batch_size, self.hidden_dim), cell_init.eval()
      )
  
      for _ in range(time_step):
          output, hidden, cell = self.sess.run(
              [output_t, h_t, c_t],
              feed_dict={
                  _step_input_var: obs_input,
                  self._step_hidden_var: hidden,
                  self._step_cell_var: cell,
              },
          )
          assert output.shape == (self.batch_size, output_dim)
          assert hidden.shape == (self.batch_size, self.hidden_dim)
          assert cell.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_lstm.py
imports: import numpy
  import pytest
  import tensorflow
  from garage.tf.models.lstm import lstm
  from tests.fixtures import TfGraphTestCase
  from tests.helpers import recurrent_step_lstm
label_desc: Use image node_id 2 for calling the TestLSTM obj's underlying member method code with example usage: obj.test_output_shapes(time_step, input_dim, output_dim, hidden_init, cell_init) without return types