docstr: Tests if the train step updates the parameters in train state.
docstr_tok: ["Tests","if","the","train","step","updates","the","parameters","in","train","state","."]
code: def test_train_step_updates_parameters(self):
      """Tests if the train step updates the parameters in train state."""
      # Create model and a state that contains the parameters.
      config = default_config.get_config()
      config.vocab_size = 13
      rng = jax.random.key(config.seed)
      model = train.model_from_config(config)
      state = train.create_train_state(rng, config, model)
  
      token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)
      lengths = np.array([2, 3], dtype=np.int32)
      labels = np.zeros_like(lengths)
      batch = {
          "token_ids": token_ids,
          "length": lengths,
          "label": labels,
      }
      rngs = {"dropout": rng}
      train_step_fn = jax.jit(train.train_step)
      new_state, metrics = train_step_fn(state, batch, rngs)
      self.assertIsInstance(new_state, train.TrainState)
      self.assertIsInstance(metrics, train.Metrics)
      old_param_values = jax.tree_util.tree_leaves(state.params)
      new_param_values = jax.tree_util.tree_leaves(new_state.params)
      for old_array, new_array in zip(
          old_param_values, new_param_values
      ):
          # Make sure parameters were updated.
          self.assertFalse(np.allclose(old_array, new_array))
  
file_path: flax/examples/sst2/train_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.test_util
  import numpy
  from configs import default
  import train
label_desc: Use image node_id 1 for calling the TrainTest obj's underlying member method code with example usage: obj.test_train_step_updates_parameters() without return types