# 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):


from typing import NamedTuple, Union
from optax._src import base
import chex
import jax
import jax.numpy as jnp

  The inner transformation processes real parameters and updates, and the
  pairs of transformed real updates are merged into complex updates.

  Parameters and updates that are real before splitting are passed through
  unmodified.

  Args:
    inner: The inner transformation.

  Returns:
    An `optax.GradientTransformation`.
  """

  def init_fn(params):
    params = jax.tree_map(_complex_to_real_pair, params)
    inner_state = inner.init(params)
    return SplitRealAndImaginaryState(inner_state)

  def update_fn(updates, state, params=None):
    inner_state = state.inner_state
    updates = jax.tree_map(_complex_to_real_pair, updates)
    params = jax.tree_map(_complex_to_real_pair, params)
    updates, inner_state = inner.update(updates, inner_state, params)
    updates = jax.tree_map(
        _real_pair_to_complex,
        updates,
        is_leaf=lambda x: isinstance(x, SplitRealAndImaginaryArrays))
    return updates, SplitRealAndImaginaryState(inner_state)


# 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/experimental/complex_valued.py
from typing import NamedTuple, Union
from optax._src import base
import chex
import jax
import jax.numpy as jnp


  The inner transformation processes real parameters and updates, and the
  pairs of transformed real updates are merged into complex updates.

  Parameters and updates that are real before splitting are passed through
  unmodified.

  Args:
    inner: The inner transformation.

  Returns:
    An `optax.GradientTransformation`.
  """

  def init_fn(params):
    params = jax.tree_map(_complex_to_real_pair, params)
    inner_state = inner.init(params)
    return SplitRealAndImaginaryState(inner_state)

  def update_fn(updates, state, params=None):
    inner_state = state.inner_state
    updates = jax.tree_map(_complex_to_real_pair, updates)
    params = jax.tree_map(_complex_to_real_pair, params)
    updates, inner_state = inner.update(updates, inner_state, params)
    updates = jax.tree_map(
        _real_pair_to_complex,
        updates,
        is_leaf=lambda x: isinstance(x, SplitRealAndImaginaryArrays))
    return updates, SplitRealAndImaginaryState(inner_state)


  return base.GradientTransformation(init_fn, update_fn)