# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/utils.py
# def linf_project_fn(epsilon: float, bounds: Tuple[float, float]) -> ProjectFn:
#   def project_fn(x, origin_x):
# def bounded_initialize_fn(
#     bounds: Optional[Tuple[chex.Array, chex.Array]] = None,) -> InitializeFn:
#     def _initialize_fn(rng, x):
# def noop_initialize_fn() -> InitializeFn:
#   def _initialize_fn(rng, x):


import collections
import chex
import jax
import jax.numpy as jnp
from typing import Callable, Optional, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils
    return _State(jnp.array(0, dtype=jnp.int32), rng, ())

  def minimize(self, x: chex.Array,
               state: _State) -> Tuple[chex.Array, chex.Array, _State]:
    """Performs a single minimization step."""
    lr = jnp.array(self._learning_rate)
    g, loss = grad_fn(self._loss_fn)(x)
    if g is None:
      raise ValueError('loss_fn does not depend on input.')
    g = jnp.sign(g)
    g, s = self._update(lr, g, state.state)
    new_state = _State(state.iteration + 1, state.rng, s)
    return x - g, loss, new_state

  def _update(
      self,
      learning_rate: chex.Numeric,
      gradients: chex.Array,
      state: chex.Array,
  ) -> Tuple[chex.Array, chex.Array]:
    return learning_rate.astype(gradients.dtype) * gradients, state


class PGD:
  """Uses the above defined optimizers to minimize and loss function."""

  def __init__(
      self,
      optimizer,
      num_steps: int,

# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/utils.py
# def linf_project_fn(epsilon: float, bounds: Tuple[float, float]) -> ProjectFn:
#   def project_fn(x, origin_x):
# def bounded_initialize_fn(
#     bounds: Optional[Tuple[chex.Array, chex.Array]] = None,) -> InitializeFn:
#     def _initialize_fn(rng, x):
# def noop_initialize_fn() -> InitializeFn:
#   def _initialize_fn(rng, x):


# Path: jax_verify/extensions/functional_lagrangian/inner_solvers/pga/optimizer.py
import collections
import chex
import jax
import jax.numpy as jnp
from typing import Callable, Optional, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils

    return _State(jnp.array(0, dtype=jnp.int32), rng, ())

  def minimize(self, x: chex.Array,
               state: _State) -> Tuple[chex.Array, chex.Array, _State]:
    """Performs a single minimization step."""
    lr = jnp.array(self._learning_rate)
    g, loss = grad_fn(self._loss_fn)(x)
    if g is None:
      raise ValueError('loss_fn does not depend on input.')
    g = jnp.sign(g)
    g, s = self._update(lr, g, state.state)
    new_state = _State(state.iteration + 1, state.rng, s)
    return x - g, loss, new_state

  def _update(
      self,
      learning_rate: chex.Numeric,
      gradients: chex.Array,
      state: chex.Array,
  ) -> Tuple[chex.Array, chex.Array]:
    return learning_rate.astype(gradients.dtype) * gradients, state


class PGD:
  """Uses the above defined optimizers to minimize and loss function."""

  def __init__(
      self,
      optimizer,
      num_steps: int,

      initialize_fn: Optional[utils.InitializeFn] = None,