def METHOD_NAME(self, time_major, go_backwards):
    input_shape = 10
    rnn_state_size = 8
    timestep = 4
    batch = 100
    x_train = np.random.random((batch, timestep, input_shape))
    def build_model(layer_cls):
        inputs = keras.layers.Input(
            shape=[timestep, input_shape], dtype=tf.float32
        )
        layer = layer_cls(
            rnn_state_size,
            recurrent_activation="sigmoid",
            time_major=time_major,
            return_sequences=True,
            go_backwards=go_backwards,
            reset_after=True,
        )
        if time_major:
            converted_input = keras.layers.Lambda(
                lambda t: tf.transpose(t, [1, 0, 2])
            )(inputs)
            outputs = layer(converted_input)
            outputs = keras.layers.Lambda(
                lambda t: tf.transpose(t, [1, 0, 2])
            )(outputs)
        else:
            outputs = layer(inputs)
        return keras.models.Model(inputs, outputs)
    gru_model = build_model(gru_v1.GRU)
    y_ref = gru_model.predict(x_train)
    weights = gru_model.get_weights()
    gru_v2_model = build_model(gru.GRU)
    gru_v2_model.set_weights(weights)
    y = gru_v2_model.predict(x_train)
    self.assertAllClose(y, y_ref)