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


from typing import Callable, List, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils
import chex
import jax
import jax.numpy as jnp
    b = jnp.where(b < 0, large_step, b)
    idx = jnp.minimum(jnp.argmin(b), num_values - 1)
    return values[idx].astype(dtype)

  return _get


class Square:
  """Performs a blackbox optimization as in https://arxiv.org/pdf/1912.00049."""

  def __init__(
      self,
      num_steps: int,
      epsilon: chex.Numeric,
      initialize_fn: utils.InitializeFn,
      bounds: Tuple[chex.ArrayTree, chex.ArrayTree],
  ):
    """Creates a Square attack."""
    self._num_steps = num_steps
    self._initialize_fn = initialize_fn
    self._project_fn = utils.linf_project_fn(epsilon=epsilon, bounds=bounds)
    self._epsilon = epsilon
    self._p_init = p = .8
    self._p_schedule = _schedule([
        p, p / 2, p / 4, p / 4, p / 8, p / 16, p / 32, p / 64, p / 128, p / 256,
        p / 512
    ], [10, 50, 200, 500, 1000, 2000, 4000, 6000, 8000])

  def __call__(
      self,

# 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/square.py
from typing import Callable, List, Tuple
from jax_verify.extensions.functional_lagrangian.inner_solvers.pga import utils
import chex
import jax
import jax.numpy as jnp

    b = jnp.where(b < 0, large_step, b)
    idx = jnp.minimum(jnp.argmin(b), num_values - 1)
    return values[idx].astype(dtype)

  return _get


class Square:
  """Performs a blackbox optimization as in https://arxiv.org/pdf/1912.00049."""

  def __init__(
      self,
      num_steps: int,
      epsilon: chex.Numeric,
      initialize_fn: utils.InitializeFn,
      bounds: Tuple[chex.ArrayTree, chex.ArrayTree],
  ):
    """Creates a Square attack."""
    self._num_steps = num_steps
    self._initialize_fn = initialize_fn
    self._project_fn = utils.linf_project_fn(epsilon=epsilon, bounds=bounds)
    self._epsilon = epsilon
    self._p_init = p = .8
    self._p_schedule = _schedule([
        p, p / 2, p / 4, p / 4, p / 8, p / 16, p / 32, p / 64, p / 128, p / 256,
        p / 512
    ], [10, 50, 200, 500, 1000, 2000, 4000, 6000, 8000])

  def __call__(
      self,

      loss_fn: utils.LossFn,