# Path: optax/_src/linear_algebra.py
# def global_norm(updates: base.Updates) -> base.Updates:
# def power_iteration(
#     matrix,
#     num_iters=100,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _iter_condition(state):
#   def _iter_body(state):
# def matrix_inverse_pth_root(
#     matrix,
#     p,
#     num_iters=100,
#     ridge_epsilon=1e-6,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _unrolled_mat_pow_1(mat_m):
#   def _unrolled_mat_pow_2(mat_m):
#   def _unrolled_mat_pow_4(mat_m):
#   def _unrolled_mat_pow_8(mat_m):
#   def mat_power(mat_m, p):
#   def _iter_condition(state):
#   def _iter_body(state):
# 
# 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 chex
import jax
import jax.numpy as jnp
import jax.scipy.stats.norm as multivariate_normal
from optax._src import linear_algebra
from optax._src import numerics


def _scale_gradient_bwd(scale, g):
  return (jax.tree_map(lambda g_: g_ * scale, g), None)


_scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)


def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
  """Scales gradients for the backwards pass.

  Args:
    inputs: A nested array.
    scale: The scale factor for the gradient on the backwards pass.

  Returns:
    An array of the same structure as `inputs`, with scaled backward gradient.
  """
  # Special case scales of 1. and 0. for more efficiency.
  if scale == 1.:
    return inputs
  elif scale == 0.:
    return jax.lax.stop_gradient(inputs)
  else:
    return _scale_gradient(inputs, scale)


# TODO(b/183800387): remove legacy aliases.
safe_norm = numerics.safe_norm

# Path: optax/_src/linear_algebra.py
# def global_norm(updates: base.Updates) -> base.Updates:
# def power_iteration(
#     matrix,
#     num_iters=100,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _iter_condition(state):
#   def _iter_body(state):
# def matrix_inverse_pth_root(
#     matrix,
#     p,
#     num_iters=100,
#     ridge_epsilon=1e-6,
#     error_tolerance=1e-6,
#     precision=lax.Precision.HIGHEST):
#   def _unrolled_mat_pow_1(mat_m):
#   def _unrolled_mat_pow_2(mat_m):
#   def _unrolled_mat_pow_4(mat_m):
#   def _unrolled_mat_pow_8(mat_m):
#   def mat_power(mat_m, p):
#   def _iter_condition(state):
#   def _iter_body(state):
# 
# 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/utils.py
import chex
import jax
import jax.numpy as jnp
import jax.scipy.stats.norm as multivariate_normal
from optax._src import linear_algebra
from optax._src import numerics



def _scale_gradient_bwd(scale, g):
  return (jax.tree_map(lambda g_: g_ * scale, g), None)


_scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)


def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
  """Scales gradients for the backwards pass.

  Args:
    inputs: A nested array.
    scale: The scale factor for the gradient on the backwards pass.

  Returns:
    An array of the same structure as `inputs`, with scaled backward gradient.
  """
  # Special case scales of 1. and 0. for more efficiency.
  if scale == 1.:
    return inputs
  elif scale == 0.:
    return jax.lax.stop_gradient(inputs)
  else:
    return _scale_gradient(inputs, scale)


# TODO(b/183800387): remove legacy aliases.
safe_norm = numerics.safe_norm

safe_int32_increment = numerics.safe_int32_increment