class_docstr: Class for testing optimizers.py.
class_docstr_tok: ["Class","for","testing","optimizers.py","."]
docstr: Optimizer state (sgd with momentum) is not updated for frozen params.
docstr_tok: ["Optimizer","state","(","sgd","with","momentum",")","is","not","updated","for","frozen","params","."]
code: def test_freeze_params_optimizer_state_frozen(
      self, freeze_params_reg_exp, frozen_params
  ):
      """Optimizer state (sgd with momentum) is not updated for frozen params."""
      optimizer_config = ml_collections.ConfigDict()
      optimizer_config.optimizer = "sgd"
      optimizer_config.momentum = 0.9
      optimizer_config.freeze_params_reg_exp = freeze_params_reg_exp
      optimizer = optimizers.get_optimizer(
          optimizer_config, self.lr, self.params
      )
      optimizer_state = optimizer.init(self.params)
  
      optimizer_dict_state = flax.traverse_util.flatten_dict(
          optimizer_state[0].inner_state[0].trace, sep="/"
      )
  
      for param_name in self.expected_sgd_updates:
          if param_name in frozen_params:
              self.assertIsInstance(
                  optimizer_dict_state[param_name], optax.MaskedNode
              )
          else:
              self.assertAllEqual(
                  optimizer_dict_state[param_name],
                  jnp.zeros_like(self.expected_sgd_updates[param_name]),
              )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  import tensorflow
label_desc: Use image node_id 6 for calling the OptimizersTest obj's underlying member method code with example usage: obj.test_freeze_params_optimizer_state_frozen(freeze_params_reg_exp, frozen_params) without return types