@pytest.mark.parametrize('time_step, input_dim, output_dim', [(1, 1, 1), (1, 1, 3), (1, 3, 1), (3, 1, 1), (3, 3, 1), (3, 3, 3)])
def test_output_value_trainable_hidden_and_cell(self, time_step, input_dim, output_dim):
    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_trainable=True, output_nonlinearity_layer=output_nonlinearity)
    self.sess.run(tf.compat.v1.global_variables_initializer())
    (outputs_t, output_t, h_t, hidden_init) = self.gru
    hidden = np.full((self.batch_size, self.hidden_dim), hidden_init.eval())
    (_, hidden) = self.sess.run([output_t, h_t], feed_dict={step_input_var: obs_input, self.step_hidden_var: hidden})
    with tf.compat.v1.variable_scope('GRU/gru', reuse=True):
        hidden_init_var = tf.compat.v1.get_variable(name='initial_hidden')
        assert (hidden_init_var in tf.compat.v1.trainable_variables())
    full_output1 = self.sess.run(outputs_t, feed_dict={input_var: obs_inputs})
    hidden2 = np.full((self.batch_size, self.hidden_dim), hidden_init.eval())
    stack_hidden = None
    for i in range(time_step):
        hidden2 = recurrent_step_gru(input_val=obs_inputs[:, i, :], num_units=self.hidden_dim, step_hidden=hidden2, w_x_init=1.0, w_h_init=1.0, b_init=0.0, nonlinearity=np.tanh, gate_nonlinearity=(lambda x: (1.0 / (1.0 + np.exp((- x))))))
        if (stack_hidden is None):
            stack_hidden = hidden2[:, np.newaxis, :]
        else:
            stack_hidden = np.concatenate((stack_hidden, hidden2[:, np.newaxis, :]), axis=1)
    output_nonlinearity = np.full((np.prod(hidden2.shape[1:]), output_dim), 1.0)
    full_output2 = np.matmul(stack_hidden, output_nonlinearity)
    assert np.allclose(full_output1, full_output2)
@pytest.mark.parametrize('time_step, input_dim, output_dim', [(1, 1, 1), (1, 1, 3), (1, 3, 1), (3, 1, 1), (3, 3, 1), (3, 3, 3)])
def test_output_value_trainable_hidden_and_cell(self, time_step, input_dim, output_dim):
    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_trainable=True, output_nonlinearity_layer=output_nonlinearity)
    self.sess.run(tf.compat.v1.global_variables_initializer())
    (outputs_t, output_t, h_t, hidden_init) = self.gru
    hidden = np.full((self.batch_size, self.hidden_dim), hidden_init.eval())
    (_, hidden) = self.sess.run([output_t, h_t], feed_dict={step_input_var: obs_input, self.step_hidden_var: hidden})
    with tf.compat.v1.variable_scope('GRU/gru', reuse=True):
        hidden_init_var = tf.compat.v1.get_variable(name='initial_hidden')
        assert (hidden_init_var in tf.compat.v1.trainable_variables())
    full_output1 = self.sess.run(outputs_t, feed_dict={input_var: obs_inputs})
    hidden2 = np.full((self.batch_size, self.hidden_dim), hidden_init.eval())
    stack_hidden = None
    for i in range(time_step):
        hidden2 = recurrent_step_gru(input_val=obs_inputs[:, i, :], num_units=self.hidden_dim, step_hidden=hidden2, w_x_init=1.0, w_h_init=1.0, b_init=0.0, nonlinearity=np.tanh, gate_nonlinearity=(lambda x: (1.0 / (1.0 + np.exp((- x))))))
        if (stack_hidden is None):
            stack_hidden = hidden2[:, np.newaxis, :]
        else:
            stack_hidden = np.concatenate((stack_hidden, hidden2[:, np.newaxis, :]), axis=1)
    output_nonlinearity = np.full((np.prod(hidden2.shape[1:]), output_dim), 1.0)
    full_output2 = np.matmul(stack_hidden, output_nonlinearity)
    assert np.allclose(full_output1, full_output2)