class_docstr: Class for testing optimizers.py.
class_docstr_tok: ["Class","for","testing","optimizers.py","."]
docstr: Creates parameters and gradient function.
docstr_tok: ["Creates","parameters","and","gradient","function","."]
code: def setUp(self):
      """Creates parameters and gradient function."""
      super().setUp()
      rng = jax.random.PRNGKey(0)
      model = MLP()
      config = ml_collections.ConfigDict()
  
      self.params, _, _, _ = train_utils.initialize_model(
          model_def=model,
          input_spec=[((1, 2), jnp.float32)],
          config=config,
          rngs=rng,
      )
  
      def training_loss_fn(params, label):
          prediction = model.apply(
              variables={"params": params}, x=jnp.array([[1.0, 1.0]])
          )
          return jnp.mean(jnp.square(prediction - label))
  
      self.compute_gradient_fn = jax.value_and_grad(training_loss_fn)
      self.label_causing_no_loss = jnp.array([10.0])
      self.label_causing_loss = jnp.array([9.0])
  
      self.lr = 0.1
      self.expected_sgd_updates = {  # For LR = 0.1
          DENSE_0_BIAS: jnp.array([-0.2, -0.2, -0.2]),
          DENSE_0_KERNEL: jnp.array(
              [[-0.2, -0.2, -0.2], [-0.2, -0.2, -0.2]]
          ),
          DENSE_1_BIAS: jnp.array([-0.2]),
          DENSE_1_KERNEL: jnp.array([[-0.6], [-0.6], [-0.6]]),
      }
  
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 1 for calling the OptimizersTest obj's underlying member method code with example usage: obj.setUp() and returns: jnp