class_docstr: Class for testing optimizers.py.
class_docstr_tok: ["Class","for","testing","optimizers.py","."]
docstr: Test freezing of parameters.
docstr_tok: ["Test","freezing","of","parameters","."]
code: def test_sgd_freeze_params(
      self, freeze_params_reg_exp, frozen_params
  ):
      """Test freezing of parameters."""
      optimizer_config = ml_collections.ConfigDict()
      optimizer_config.optimizer = "sgd"
      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)
  
      _, grad = self.compute_gradient_fn(
          self.params, self.label_causing_loss
      )
      updates, _ = optimizer.update(grad, optimizer_state, self.params)
      updates = flax.traverse_util.flatten_dict(updates, sep="/")
  
      for param_name in self.expected_sgd_updates:
          if param_name in frozen_params:
              self.assertAllClose(
                  updates[param_name],
                  jnp.zeros_like(self.expected_sgd_updates[param_name]),
              )
          else:
              self.assertAllClose(
                  updates[param_name],
                  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 5 for calling the OptimizersTest obj's underlying member method code with example usage: obj.test_sgd_freeze_params(freeze_params_reg_exp, frozen_params) without return types