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


import functools
import inspect
import chex
import jax
import jax.numpy as jnp
from typing import Callable, Dict, Union, NamedTuple, Optional, Iterable, Sequence
from absl import logging
from optax._src import base
from optax._src import numerics
    raise ValueError('`interpolate_type` must be either \'cos\' or \'linear\'')

  if boundaries_and_scales:
    boundaries, scales = zip(*sorted(boundaries_and_scales.items()))
    if not all(scale >= 0. for scale in scales):
      raise ValueError(
          '`piecewise_interpolate_schedule` expects non-negative scale factors')
  else:
    boundaries, scales = (), ()

  bounds = jnp.stack((0,) + boundaries)
  values = jnp.cumprod(jnp.stack((init_value,) + scales))
  interval_sizes = (bounds[1:] - bounds[:-1])

  def schedule(count):
    indicator = (bounds[:-1] <= count) & (count < bounds[1:])
    pct = (count - bounds[:-1]) / interval_sizes
    interp_vals = interpolate_fn(values[:-1], values[1:], pct)
    return indicator.dot(interp_vals) + (bounds[-1] <= count) * values[-1]

  return schedule


def linear_onecycle_schedule(
    transition_steps: int,
    peak_value: float,
    pct_start: float = 0.3,
    pct_final: float = 0.85,
    div_factor: float = 25.0,
    final_div_factor: float = 1e4

# 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/schedule.py
import functools
import inspect
import chex
import jax
import jax.numpy as jnp
from typing import Callable, Dict, Union, NamedTuple, Optional, Iterable, Sequence
from absl import logging
from optax._src import base
from optax._src import numerics

    raise ValueError('`interpolate_type` must be either \'cos\' or \'linear\'')

  if boundaries_and_scales:
    boundaries, scales = zip(*sorted(boundaries_and_scales.items()))
    if not all(scale >= 0. for scale in scales):
      raise ValueError(
          '`piecewise_interpolate_schedule` expects non-negative scale factors')
  else:
    boundaries, scales = (), ()

  bounds = jnp.stack((0,) + boundaries)
  values = jnp.cumprod(jnp.stack((init_value,) + scales))
  interval_sizes = (bounds[1:] - bounds[:-1])

  def schedule(count):
    indicator = (bounds[:-1] <= count) & (count < bounds[1:])
    pct = (count - bounds[:-1]) / interval_sizes
    interp_vals = interpolate_fn(values[:-1], values[1:], pct)
    return indicator.dot(interp_vals) + (bounds[-1] <= count) * values[-1]

  return schedule


def linear_onecycle_schedule(
    transition_steps: int,
    peak_value: float,
    pct_start: float = 0.3,
    pct_final: float = 0.85,
    div_factor: float = 25.0,
    final_div_factor: float = 1e4

) -> base.Schedule: