class_docstr: Class for testing optimizers.py.
class_docstr_tok: ["Class","for","testing","optimizers.py","."]
docstr: Test SGD with weight decay.
docstr_tok: ["Test","SGD","with","weight","decay","."]
code: def test_sgd_with_weight_decay(self, weight_decay):
      """Test SGD with weight decay."""
      optimizer_config = ml_collections.ConfigDict()
      optimizer_config.optimizer = "sgd"
      optimizer_config.weight_decay = weight_decay
      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:
          self.assertAllClose(
              updates[param_name],
              self.expected_sgd_updates[param_name]
              - weight_decay * self.lr,
          )
  
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 3 for calling the OptimizersTest obj's underlying member method code with example usage: obj.test_sgd_with_weight_decay(weight_decay) without return types