# Path: jax_verify/extensions/functional_lagrangian/verify_utils.py
# class AbstractParams(abc.ABC):
# class FCParams(AbstractParams):
# class ConvParams(AbstractParams):
# class SpecType(enum.Enum):
# class Distribution(enum.Enum):
# class NetworkType(enum.Enum):
# class InnerVerifInstance:
#   def __call__(self, inputs: Tensor) -> Tensor:
#   def params(self):
#   def has_bounds(self):
#   def params(self):
#   def params(self):
#   def same_lagrangian_form_pre_post(self) -> bool:
#   UNCERTAINTY = 'uncertainty'
#   ADVERSARIAL = 'adversarial'
#   ADVERSARIAL_SOFTMAX = 'adversarial_softmax'
#   PROBABILITY_THRESHOLD = 'probability_threshold'
#   GAUSSIAN = 'gaussian'
#   BERNOULLI = 'bernoulli'
#   DETERMINISTIC = 'deterministic'
#   STOCHASTIC = 'stochastic'
# 
# Path: jax_verify/extensions/sdp_verify/utils.py
# class SdpDualVerifInstance(_SdpDualVerifInstance):
# class VerifInstanceTypes(enum.Enum):
# class DualVarTypes(enum.Enum):
# def mlp_layer_sizes(params):
# def nn_layer_sizes(params):
# def layer_sizes_from_bounds(bounds):
# def predict_mlp(params, inputs):
# def fwd(inputs, layer_params):
# def predict_cnn(params, inputs, include_preactivations=False):
# def get_network_activs(params, x):
# def get_layer_params(fun_to_extract, example_input):
#   def _get_const_input_arg(eqn):
# def boundprop(params, bounds_in):
# def init_bound(x, epsilon, input_bounds=(0., 1.), add_batch_dim=True):
# def ibp_bound_elided(verif_instance):
# def ibp_bound_nonelided(verif_instance):
# def make_relu_robust_verif_instance(
#     params, bounds=None, target_label=1, label=2, input_bounds=None):
# def make_relu_robust_verif_instance_elided(
#     params, bounds=None, input_bounds=None):
# def output_size_and_verif_type(params, bounds):
# def elide_params(params, label, target_label, op_size):
# def preprocess_cifar(image, inception_preprocess=False, perturbation=False):
# def preprocessed_cifar_eps_and_input_bounds(
#     shape=(32, 32, 3), epsilon=2/255, inception_preprocess=False):
# def adv_objective(model_fn, x, label, target_label):
# def fgsm_single(model_fn, x, label, target_label, epsilon, num_steps,
#                 step_size, input_bounds=(0., 1.)):
# def untargeted_margin_loss(logits, labels):
# def pgd_default(model_fn, x, label, epsilon, num_steps, step_size,
#                 input_bounds=(0., 1.)):
# def pgd(adv_loss, x_init, epsilon, num_steps, step_size, input_bounds=(0., 1.)):
# def scale_by_variable_opt(multipliers):
#   def init_fn(params):
#   def update_fn(updates, _, params=None):
# def flatten(pytree, backend=np):
# def unflatten_like(a, pytree):
# def structure_like(tree1, tree2):
#   MLP_ELIDED = 'mlp_elided'
#   CNN_ELIDED = 'cnn_elided'
#   EQUALITY = 'equality'
#   INEQUALITY = 'inequality'


import dataclasses
import jax
import jax.numpy as jnp
import jax_verify
import optax
from typing import Callable, Union
from jax_verify.extensions.functional_lagrangian import verify_utils
from jax_verify.extensions.sdp_verify import utils as sdp_utils
    learning_rate: float,
    num_samples: int = 1,
) -> float:
  """Adversarial attack on uncertainty spec (with parameter sampling)."""
  l = jnp.clip(data_spec.input-data_spec.epsilon,
               data_spec.input_bounds[0], data_spec.input_bounds[1])
  u = jnp.clip(data_spec.input+data_spec.epsilon,
               data_spec.input_bounds[0], data_spec.input_bounds[1])
  projection_fn = lambda x: jnp.clip(x, l, u)

  forward_fn = make_forward(params, num_samples)

  def max_objective_fn_uncertainty(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    return logits[data_spec.target_label]

  def max_objective_fn_adversarial(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    return logits[data_spec.target_label] - logits[data_spec.true_label]

  def max_objective_fn_adversarial_softmax(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    probs = jax.nn.softmax(logits, axis=-1)
    return probs[data_spec.target_label] - probs[data_spec.true_label]

  if (spec_type in (verify_utils.SpecType.UNCERTAINTY,
                    verify_utils.SpecType.PROBABILITY_THRESHOLD)):
    max_objective_fn = max_objective_fn_uncertainty
  elif spec_type == verify_utils.SpecType.ADVERSARIAL:
    max_objective_fn = max_objective_fn_adversarial

# Path: jax_verify/extensions/functional_lagrangian/verify_utils.py
# class AbstractParams(abc.ABC):
# class FCParams(AbstractParams):
# class ConvParams(AbstractParams):
# class SpecType(enum.Enum):
# class Distribution(enum.Enum):
# class NetworkType(enum.Enum):
# class InnerVerifInstance:
#   def __call__(self, inputs: Tensor) -> Tensor:
#   def params(self):
#   def has_bounds(self):
#   def params(self):
#   def params(self):
#   def same_lagrangian_form_pre_post(self) -> bool:
#   UNCERTAINTY = 'uncertainty'
#   ADVERSARIAL = 'adversarial'
#   ADVERSARIAL_SOFTMAX = 'adversarial_softmax'
#   PROBABILITY_THRESHOLD = 'probability_threshold'
#   GAUSSIAN = 'gaussian'
#   BERNOULLI = 'bernoulli'
#   DETERMINISTIC = 'deterministic'
#   STOCHASTIC = 'stochastic'
# 
# Path: jax_verify/extensions/sdp_verify/utils.py
# class SdpDualVerifInstance(_SdpDualVerifInstance):
# class VerifInstanceTypes(enum.Enum):
# class DualVarTypes(enum.Enum):
# def mlp_layer_sizes(params):
# def nn_layer_sizes(params):
# def layer_sizes_from_bounds(bounds):
# def predict_mlp(params, inputs):
# def fwd(inputs, layer_params):
# def predict_cnn(params, inputs, include_preactivations=False):
# def get_network_activs(params, x):
# def get_layer_params(fun_to_extract, example_input):
#   def _get_const_input_arg(eqn):
# def boundprop(params, bounds_in):
# def init_bound(x, epsilon, input_bounds=(0., 1.), add_batch_dim=True):
# def ibp_bound_elided(verif_instance):
# def ibp_bound_nonelided(verif_instance):
# def make_relu_robust_verif_instance(
#     params, bounds=None, target_label=1, label=2, input_bounds=None):
# def make_relu_robust_verif_instance_elided(
#     params, bounds=None, input_bounds=None):
# def output_size_and_verif_type(params, bounds):
# def elide_params(params, label, target_label, op_size):
# def preprocess_cifar(image, inception_preprocess=False, perturbation=False):
# def preprocessed_cifar_eps_and_input_bounds(
#     shape=(32, 32, 3), epsilon=2/255, inception_preprocess=False):
# def adv_objective(model_fn, x, label, target_label):
# def fgsm_single(model_fn, x, label, target_label, epsilon, num_steps,
#                 step_size, input_bounds=(0., 1.)):
# def untargeted_margin_loss(logits, labels):
# def pgd_default(model_fn, x, label, epsilon, num_steps, step_size,
#                 input_bounds=(0., 1.)):
# def pgd(adv_loss, x_init, epsilon, num_steps, step_size, input_bounds=(0., 1.)):
# def scale_by_variable_opt(multipliers):
#   def init_fn(params):
#   def update_fn(updates, _, params=None):
# def flatten(pytree, backend=np):
# def unflatten_like(a, pytree):
# def structure_like(tree1, tree2):
#   MLP_ELIDED = 'mlp_elided'
#   CNN_ELIDED = 'cnn_elided'
#   EQUALITY = 'equality'
#   INEQUALITY = 'inequality'


# Path: jax_verify/extensions/functional_lagrangian/attacks.py
import dataclasses
import jax
import jax.numpy as jnp
import jax_verify
import optax
from typing import Callable, Union
from jax_verify.extensions.functional_lagrangian import verify_utils
from jax_verify.extensions.sdp_verify import utils as sdp_utils

    learning_rate: float,
    num_samples: int = 1,
) -> float:
  """Adversarial attack on uncertainty spec (with parameter sampling)."""
  l = jnp.clip(data_spec.input-data_spec.epsilon,
               data_spec.input_bounds[0], data_spec.input_bounds[1])
  u = jnp.clip(data_spec.input+data_spec.epsilon,
               data_spec.input_bounds[0], data_spec.input_bounds[1])
  projection_fn = lambda x: jnp.clip(x, l, u)

  forward_fn = make_forward(params, num_samples)

  def max_objective_fn_uncertainty(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    return logits[data_spec.target_label]

  def max_objective_fn_adversarial(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    return logits[data_spec.target_label] - logits[data_spec.true_label]

  def max_objective_fn_adversarial_softmax(x, prng_key):
    logits = jnp.reshape(forward_fn(x, prng_key), [-1])
    probs = jax.nn.softmax(logits, axis=-1)
    return probs[data_spec.target_label] - probs[data_spec.true_label]

  if (spec_type in (verify_utils.SpecType.UNCERTAINTY,
                    verify_utils.SpecType.PROBABILITY_THRESHOLD)):
    max_objective_fn = max_objective_fn_uncertainty
  elif spec_type == verify_utils.SpecType.ADVERSARIAL:
    max_objective_fn = max_objective_fn_adversarial

  elif spec_type == verify_utils.SpecType.ADVERSARIAL_SOFTMAX: