# Path: optax/_src/update.py
# def apply_updates(params: base.Params, updates: base.Updates) -> base.Params:
# def incremental_update(
#     new_tensors: base.Params,
#     old_tensors: base.Params,
#     step_size: chex.Numeric
# ) -> base.Params:
# def periodic_update(
#     new_tensors: base.Params,
#     old_tensors: base.Params,
#     steps: chex.Array,
#     update_period: int
# ) -> base.Params:


from absl.testing import absltest
from optax._src import update
import chex
import jax
import jax.numpy as jnp
class UpdateTest(chex.TestCase):

  @chex.all_variants()
  def test_apply_updates(self):
    params = ({'a': jnp.ones((3, 2))}, jnp.ones((1,)))
    grads = jax.tree_map(lambda t: 2 * t, params)
    exp_params = jax.tree_map(lambda t: 3 * t, params)
    new_params = self.variant(update.apply_updates)(params, grads)

    chex.assert_tree_all_close(
        exp_params, new_params, atol=1e-10, rtol=1e-5)

  @chex.all_variants()
  def test_apply_updates_mixed_precision(self):
    params = (
        {'a': jnp.ones((3, 2), dtype=jnp.bfloat16)},
        jnp.ones((1,), dtype=jnp.bfloat16))
    grads = jax.tree_map(
        lambda t: (2 * t).astype(jnp.float32), params)
    new_params = self.variant(update.apply_updates)(params, grads)

    for leaf in jax.tree_leaves(new_params):
      assert leaf.dtype == jnp.bfloat16

  @chex.all_variants()
  def test_incremental_update(self):
    params_1 = ({'a': jnp.ones((3, 2))}, jnp.ones((1,)))
    params_2 = jax.tree_map(lambda t: 2 * t, params_1)
    exp_params = jax.tree_map(lambda t: 1.5 * t, params_1)
    new_params = self.variant(

# Path: optax/_src/update.py
# def apply_updates(params: base.Params, updates: base.Updates) -> base.Params:
# def incremental_update(
#     new_tensors: base.Params,
#     old_tensors: base.Params,
#     step_size: chex.Numeric
# ) -> base.Params:
# def periodic_update(
#     new_tensors: base.Params,
#     old_tensors: base.Params,
#     steps: chex.Array,
#     update_period: int
# ) -> base.Params:


# Path: optax/_src/update_test.py
from absl.testing import absltest
from optax._src import update
import chex
import jax
import jax.numpy as jnp

class UpdateTest(chex.TestCase):

  @chex.all_variants()
  def test_apply_updates(self):
    params = ({'a': jnp.ones((3, 2))}, jnp.ones((1,)))
    grads = jax.tree_map(lambda t: 2 * t, params)
    exp_params = jax.tree_map(lambda t: 3 * t, params)
    new_params = self.variant(update.apply_updates)(params, grads)

    chex.assert_tree_all_close(
        exp_params, new_params, atol=1e-10, rtol=1e-5)

  @chex.all_variants()
  def test_apply_updates_mixed_precision(self):
    params = (
        {'a': jnp.ones((3, 2), dtype=jnp.bfloat16)},
        jnp.ones((1,), dtype=jnp.bfloat16))
    grads = jax.tree_map(
        lambda t: (2 * t).astype(jnp.float32), params)
    new_params = self.variant(update.apply_updates)(params, grads)

    for leaf in jax.tree_leaves(new_params):
      assert leaf.dtype == jnp.bfloat16

  @chex.all_variants()
  def test_incremental_update(self):
    params_1 = ({'a': jnp.ones((3, 2))}, jnp.ones((1,)))
    params_2 = jax.tree_map(lambda t: 2 * t, params_1)
    exp_params = jax.tree_map(lambda t: 1.5 * t, params_1)
    new_params = self.variant(

        update.incremental_update)(params_2, params_1, 0.5)