# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# 
# 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
# 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():


import dataclasses
import chex
import jax
import jax.numpy as jnp
import numpy as np
from typing import NamedTuple, Optional, Tuple, Callable
from optax._src import base
from optax._src import numerics
from optax._src import utils
        v_col=jax.tree_map(lambda o: o.v_col, result_tree),
        v=jax.tree_map(lambda o: o.v, result_tree))

  def init_fn(params):
    """Initialise the optimiser's state."""

    def _init(param):
      shape = param.shape
      factored_dims = _factored_dims(shape, factored, min_dim_size_to_factor)
      if factored_dims is not None:
        d1, d0 = factored_dims
        vr_shape = np.delete(shape, d0)
        vc_shape = np.delete(shape, d1)
        return _UpdateResult(
            update=jnp.zeros((1,)),
            v_row=jnp.zeros(vr_shape),
            v_col=jnp.zeros(vc_shape),
            v=jnp.zeros((1,)))
      else:
        return _UpdateResult(
            update=jnp.zeros((1,)),
            v_row=jnp.zeros((1,)),
            v_col=jnp.zeros((1,)),
            v=jnp.zeros(param.shape))

    return _to_state(jnp.zeros([], jnp.int32), jax.tree_map(_init, params))

  def update_fn(grads, state, params):
    """Apply gradient transformation."""
    if params is None:

# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# 
# 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
# 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/factorized.py
import dataclasses
import chex
import jax
import jax.numpy as jnp
import numpy as np
from typing import NamedTuple, Optional, Tuple, Callable
from optax._src import base
from optax._src import numerics
from optax._src import utils

        v_col=jax.tree_map(lambda o: o.v_col, result_tree),
        v=jax.tree_map(lambda o: o.v, result_tree))

  def init_fn(params):
    """Initialise the optimiser's state."""

    def _init(param):
      shape = param.shape
      factored_dims = _factored_dims(shape, factored, min_dim_size_to_factor)
      if factored_dims is not None:
        d1, d0 = factored_dims
        vr_shape = np.delete(shape, d0)
        vc_shape = np.delete(shape, d1)
        return _UpdateResult(
            update=jnp.zeros((1,)),
            v_row=jnp.zeros(vr_shape),
            v_col=jnp.zeros(vc_shape),
            v=jnp.zeros((1,)))
      else:
        return _UpdateResult(
            update=jnp.zeros((1,)),
            v_row=jnp.zeros((1,)),
            v_col=jnp.zeros((1,)),
            v=jnp.zeros(param.shape))

    return _to_state(jnp.zeros([], jnp.int32), jax.tree_map(_init, params))

  def update_fn(grads, state, params):
    """Apply gradient transformation."""
    if params is None:

      raise ValueError(base.NO_PARAMS_MSG)