# Path: jax_verify/extensions/functional_lagrangian/lagrangian_form.py
# def _flatten_spatial_dims(x: Tensor) -> Tensor:
# def size_from_shape(shape: Shape) -> int:
#   def __init__(self, name):
#   def _init_params_per_sample(self, key: PRNGKey, *args) -> Params:
#   def init_params(self, key, *args, **kwargs):
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def process_params(self, lagrange_params: Params):
#   def name(self):
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = True) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = False) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
# def get_lagrangian_form(config_lagrangian_form: ConfigDict) -> LagrangianForm:
# class LagrangianForm(metaclass=abc.ABCMeta):
# class Linear(LagrangianForm):
# class LinearExp(LagrangianForm):
# 
# 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 abc
import collections
import dataclasses
import enum
import chex
import jax.numpy as jnp
import jax_verify
import ml_collections
from typing import Callable, List, Optional, Union
from jax_verify.extensions.functional_lagrangian import lagrangian_form
from jax_verify.extensions.sdp_verify import utils as sdp_utils

"""Small helper functions."""




Params = collections.namedtuple('Params', ['inner', 'outer'])
ParamsTypes = collections.namedtuple('ParamsTypes',
                                     ['inner', 'outer', 'lagrangian_form'])

DataSpec = collections.namedtuple(
    'DataSpec',
    ['input', 'true_label', 'target_label', 'epsilon', 'input_bounds'])

Array = chex.Array
ArrayTree = chex.ArrayTree
ConfigDict = ml_collections.ConfigDict
IntervalBound = jax_verify.IntervalBound
Tensor = jnp.array
LayerParams = Union['FCParams', 'ConvParams']
LagrangianForm = lagrangian_form.LagrangianForm
ModelParams = List[LayerParams]
ModelParamsElided = ModelParams


class AbstractParams(abc.ABC):
  """AbstractParams."""

  def __call__(self, inputs: Tensor) -> Tensor:
    """Forward pass on layer."""

# Path: jax_verify/extensions/functional_lagrangian/lagrangian_form.py
# def _flatten_spatial_dims(x: Tensor) -> Tensor:
# def size_from_shape(shape: Shape) -> int:
#   def __init__(self, name):
#   def _init_params_per_sample(self, key: PRNGKey, *args) -> Params:
#   def init_params(self, key, *args, **kwargs):
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def process_params(self, lagrange_params: Params):
#   def name(self):
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = True) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = False) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
# def get_lagrangian_form(config_lagrangian_form: ConfigDict) -> LagrangianForm:
# class LagrangianForm(metaclass=abc.ABCMeta):
# class Linear(LagrangianForm):
# class LinearExp(LagrangianForm):
# 
# 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/verify_utils.py
import abc
import collections
import dataclasses
import enum
import chex
import jax.numpy as jnp
import jax_verify
import ml_collections
from typing import Callable, List, Optional, Union
from jax_verify.extensions.functional_lagrangian import lagrangian_form
from jax_verify.extensions.sdp_verify import utils as sdp_utils


"""Small helper functions."""




Params = collections.namedtuple('Params', ['inner', 'outer'])
ParamsTypes = collections.namedtuple('ParamsTypes',
                                     ['inner', 'outer', 'lagrangian_form'])

DataSpec = collections.namedtuple(
    'DataSpec',
    ['input', 'true_label', 'target_label', 'epsilon', 'input_bounds'])

Array = chex.Array
ArrayTree = chex.ArrayTree
ConfigDict = ml_collections.ConfigDict
IntervalBound = jax_verify.IntervalBound
Tensor = jnp.array
LayerParams = Union['FCParams', 'ConvParams']
LagrangianForm = lagrangian_form.LagrangianForm
ModelParams = List[LayerParams]
ModelParamsElided = ModelParams


class AbstractParams(abc.ABC):
  """AbstractParams."""

  def __call__(self, inputs: Tensor) -> Tensor:
    """Forward pass on layer."""

    return sdp_utils.fwd(inputs, self.params)