# Path: optax/_src/numerics.py
# def abs_sq(x: chex.Array) -> chex.Array:
# def safe_norm(x: chex.Array,
#               min_norm: chex.Numeric,
#               ord: Optional[Union[int, float, str]] = None,  # pylint: disable=redefined-builtin
#               axis: Union[None, Tuple[int, ...], int] = None,
#               keepdims: bool = False) -> chex.Array:
# def safe_root_mean_squares(x: chex.Array, min_rms: chex.Numeric) -> chex.Array:
# def safe_int32_increment(count: chex.Numeric) -> chex.Numeric:


import functools
import itertools
import re
import chex
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from optax._src import numerics
    inc_fn = self.variant(numerics.safe_int32_increment)
    # increment small numbers correctly.
    base = int32_array(3)
    incremented = inc_fn(base)
    np.testing.assert_array_equal(incremented, int32_array(4))
    # avoid overflow when incrementing maxint.
    base = int32_array(np.iinfo(np.int32).max)
    incremented = inc_fn(base)
    np.testing.assert_array_equal(incremented, base)

  @chex.all_variants()
  @parameterized.parameters(
      itertools.filterfalse(
          _invalid_ord_axis_inputs,
          itertools.product(_ALL_ORDS, [None, 0, 1], [False, True])))
  def test_safe_norm(self, ord, axis, keepdims):  # pylint: disable=redefined-builtin
    dnorm_dx = self.variant(
        jax.jacfwd(
            functools.partial(
                numerics.safe_norm, ord=ord, axis=axis, keepdims=keepdims),
            argnums=0))
    # Test gradient is 0. in 0. when zero min norm is used.
    g = dnorm_dx(float32_array(jnp.zeros((3, 4))), float32_array(0.))
    np.testing.assert_array_equal(g, jnp.zeros_like(g))
    # Test gradient is 0. in 0. when non zero min norm is used.
    g = dnorm_dx(float32_array(jnp.zeros((3, 4))), float32_array(3.))
    np.testing.assert_array_equal(g, jnp.zeros_like(g))

  @chex.all_variants()
  def test_safe_rms(self):

# Path: optax/_src/numerics.py
# def abs_sq(x: chex.Array) -> chex.Array:
# def safe_norm(x: chex.Array,
#               min_norm: chex.Numeric,
#               ord: Optional[Union[int, float, str]] = None,  # pylint: disable=redefined-builtin
#               axis: Union[None, Tuple[int, ...], int] = None,
#               keepdims: bool = False) -> chex.Array:
# def safe_root_mean_squares(x: chex.Array, min_rms: chex.Numeric) -> chex.Array:
# def safe_int32_increment(count: chex.Numeric) -> chex.Numeric:


# Path: optax/_src/numerics_test.py
import functools
import itertools
import re
import chex
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from optax._src import numerics

    inc_fn = self.variant(numerics.safe_int32_increment)
    # increment small numbers correctly.
    base = int32_array(3)
    incremented = inc_fn(base)
    np.testing.assert_array_equal(incremented, int32_array(4))
    # avoid overflow when incrementing maxint.
    base = int32_array(np.iinfo(np.int32).max)
    incremented = inc_fn(base)
    np.testing.assert_array_equal(incremented, base)

  @chex.all_variants()
  @parameterized.parameters(
      itertools.filterfalse(
          _invalid_ord_axis_inputs,
          itertools.product(_ALL_ORDS, [None, 0, 1], [False, True])))
  def test_safe_norm(self, ord, axis, keepdims):  # pylint: disable=redefined-builtin
    dnorm_dx = self.variant(
        jax.jacfwd(
            functools.partial(
                numerics.safe_norm, ord=ord, axis=axis, keepdims=keepdims),
            argnums=0))
    # Test gradient is 0. in 0. when zero min norm is used.
    g = dnorm_dx(float32_array(jnp.zeros((3, 4))), float32_array(0.))
    np.testing.assert_array_equal(g, jnp.zeros_like(g))
    # Test gradient is 0. in 0. when non zero min norm is used.
    g = dnorm_dx(float32_array(jnp.zeros((3, 4))), float32_array(3.))
    np.testing.assert_array_equal(g, jnp.zeros_like(g))

  @chex.all_variants()
  def test_safe_rms(self):

    drms_dx = self.variant(jax.grad(numerics.safe_root_mean_squares))