# Path: optax/_src/utils.py
# def tile_second_to_last_dim(a):
# def canonicalize_dtype(dtype):
# def cast_tree(tree, dtype):
# def set_diags(a, new_diags):
# def multi_normal(loc, log_scale):
#   def __init__(self, loc, log_scale):
#   def sample(self, shape, seed):
#   def log_prob(self, x):
#   def log_scale(self):
#   def params(self):
# def _scale_gradient(inputs: chex.Array, scale: float)  -> chex.ArrayTree:
# def _scale_gradient_fwd(inputs, scale):
# def _scale_gradient_bwd(scale, g):
# def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
# class MultiNormalDiagFromLogScale():


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from optax._src import utils
import jax
# ==============================================================================
"""Tests for `utils.py`."""






class ScaleGradientTest(parameterized.TestCase):

  @parameterized.product(inputs=[-1., 0., 1.], scale=[-0.5, 0., 0.5, 1., 2.])
  @mock.patch.object(jax.lax, 'stop_gradient', wraps=jax.lax.stop_gradient)
  def test_scale_gradient(self, mock_sg, inputs, scale):

    def fn(inputs):
      outputs = utils.scale_gradient(inputs, scale)
      return outputs ** 2

    grad = jax.grad(fn)
    self.assertEqual(grad(inputs), 2 * inputs * scale)
    if scale == 0.:
      mock_sg.assert_called_once_with(inputs)
    else:
      self.assertFalse(mock_sg.called)
    self.assertEqual(fn(inputs), inputs ** 2)

  @parameterized.product(scale=[-0.5, 0., 0.5, 1., 2.])
  def test_scale_gradient_pytree(self, scale):

    def fn(inputs):

# Path: optax/_src/utils.py
# def tile_second_to_last_dim(a):
# def canonicalize_dtype(dtype):
# def cast_tree(tree, dtype):
# def set_diags(a, new_diags):
# def multi_normal(loc, log_scale):
#   def __init__(self, loc, log_scale):
#   def sample(self, shape, seed):
#   def log_prob(self, x):
#   def log_scale(self):
#   def params(self):
# def _scale_gradient(inputs: chex.Array, scale: float)  -> chex.ArrayTree:
# def _scale_gradient_fwd(inputs, scale):
# def _scale_gradient_bwd(scale, g):
# def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
# class MultiNormalDiagFromLogScale():


# Path: optax/_src/utils_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from optax._src import utils
import jax

# ==============================================================================
"""Tests for `utils.py`."""






class ScaleGradientTest(parameterized.TestCase):

  @parameterized.product(inputs=[-1., 0., 1.], scale=[-0.5, 0., 0.5, 1., 2.])
  @mock.patch.object(jax.lax, 'stop_gradient', wraps=jax.lax.stop_gradient)
  def test_scale_gradient(self, mock_sg, inputs, scale):

    def fn(inputs):
      outputs = utils.scale_gradient(inputs, scale)
      return outputs ** 2

    grad = jax.grad(fn)
    self.assertEqual(grad(inputs), 2 * inputs * scale)
    if scale == 0.:
      mock_sg.assert_called_once_with(inputs)
    else:
      self.assertFalse(mock_sg.called)
    self.assertEqual(fn(inputs), inputs ** 2)

  @parameterized.product(scale=[-0.5, 0., 0.5, 1., 2.])
  def test_scale_gradient_pytree(self, scale):

    def fn(inputs):

      outputs = utils.scale_gradient(inputs, scale)