# 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/src/utils.py
# def open_file(name, *open_args, root_dir='/tmp/jax_verify', **open_kwargs):
# def bind_nonbound_args(
#     fun: Callable[..., Tensor],
#     *all_in_args: Union[Bound, Tensor],
#     **kwargs
# ) -> Callable[..., Tensor]:
#   def tensorbound_fun(*bound_args):
# def filter_jaxverify_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
# def simple_propagation(fn):
#   def wrapper(context, *args, **kwargs):
# def batch_value_and_grad(fun, batch_dims, *args, **kwargs):
#   def nobatch_fun(*nobatch_inps):
# def objective_chunk(
#     obj_shape: Sequence[int],
#     chunk_index: int,
#     nb_parallel_nodes: int,
# ):
# def chunked_bounds(
#     bound_shape: Tuple[int, ...],
#     max_parallel_nodes: int,
#     bound_fn: Callable[[Tensor], Tuple[Tensor, Tensor]],
# ) -> Tuple[Tensor, Tensor]:
#   def bound_chunk(chunk_index: int) -> Tuple[Tensor, Tensor]:


import dataclasses
import os
import pickle
import urllib
import jax.numpy as jnp
import jax_verify
import ml_collections
import numpy as np
from typing import Any, Optional
from jax_verify.extensions.functional_lagrangian import verify_utils
from jax_verify.src import utils as jv_utils
    'cifar_vgg_32': 'models/cifar_vgg_32_dropout.pkl',
    'cifar_vgg_64': 'models/cifar_vgg_64_dropout.pkl',
})

PROBA_SAFETY_URL = (
    'https://github.com/matthewwicker/ProbabilisticSafetyforBNNs/raw/master'
    '/MNIST/concurMNIST2/MNIST_Networks')

PROBA_SAFETY_MODEL_PATHS = ml_collections.ConfigDict({
    'mnist_mlp_1_1024': 'VIMODEL_MNIST_1_1024_relu.net.npz',
    'mnist_mlp_1_128': 'VIMODEL_MNIST_1_128_relu.net.npz',
    'mnist_mlp_1_2048': 'VIMODEL_MNIST_1_2048_relu.net.npz',
    'mnist_mlp_1_256': 'VIMODEL_MNIST_1_256_relu.net.npz',
    'mnist_mlp_1_4096': 'VIMODEL_MNIST_1_4096_relu.net.npz',
    'mnist_mlp_1_512': 'VIMODEL_MNIST_1_512_relu.net.npz',
    'mnist_mlp_1_64': 'VIMODEL_MNIST_1_64_relu.net.npz',
    'mnist_mlp_2_1024': 'VIMODEL_MNIST_2_1024_relu.net.npz',
    'mnist_mlp_2_128': 'VIMODEL_MNIST_2_128_relu.net.npz',
    'mnist_mlp_2_256': 'VIMODEL_MNIST_2_256_relu.net.npz',
    'mnist_mlp_2_512': 'VIMODEL_MNIST_2_512_relu.net.npz',
    'mnist_mlp_2_64': 'VIMODEL_MNIST_2_64_relu.net.npz',
})


def _load_pickled_model(root_dir: str, model_name: str) -> ModelParams:
  model_path = getattr(INTERNAL_MODEL_PATHS, model_name.lower())
  if model_path.endswith('mnist_ceda.pkl'):
    with jv_utils.open_file(model_path, 'rb', root_dir=root_dir) as f:
      params_iterables = pickle.load(f, encoding='bytes')
  else:

# 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/src/utils.py
# def open_file(name, *open_args, root_dir='/tmp/jax_verify', **open_kwargs):
# def bind_nonbound_args(
#     fun: Callable[..., Tensor],
#     *all_in_args: Union[Bound, Tensor],
#     **kwargs
# ) -> Callable[..., Tensor]:
#   def tensorbound_fun(*bound_args):
# def filter_jaxverify_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
# def simple_propagation(fn):
#   def wrapper(context, *args, **kwargs):
# def batch_value_and_grad(fun, batch_dims, *args, **kwargs):
#   def nobatch_fun(*nobatch_inps):
# def objective_chunk(
#     obj_shape: Sequence[int],
#     chunk_index: int,
#     nb_parallel_nodes: int,
# ):
# def chunked_bounds(
#     bound_shape: Tuple[int, ...],
#     max_parallel_nodes: int,
#     bound_fn: Callable[[Tensor], Tuple[Tensor, Tensor]],
# ) -> Tuple[Tensor, Tensor]:
#   def bound_chunk(chunk_index: int) -> Tuple[Tensor, Tensor]:


# Path: jax_verify/extensions/functional_lagrangian/model.py
import dataclasses
import os
import pickle
import urllib
import jax.numpy as jnp
import jax_verify
import ml_collections
import numpy as np
from typing import Any, Optional
from jax_verify.extensions.functional_lagrangian import verify_utils
from jax_verify.src import utils as jv_utils

    'cifar_vgg_32': 'models/cifar_vgg_32_dropout.pkl',
    'cifar_vgg_64': 'models/cifar_vgg_64_dropout.pkl',
})

PROBA_SAFETY_URL = (
    'https://github.com/matthewwicker/ProbabilisticSafetyforBNNs/raw/master'
    '/MNIST/concurMNIST2/MNIST_Networks')

PROBA_SAFETY_MODEL_PATHS = ml_collections.ConfigDict({
    'mnist_mlp_1_1024': 'VIMODEL_MNIST_1_1024_relu.net.npz',
    'mnist_mlp_1_128': 'VIMODEL_MNIST_1_128_relu.net.npz',
    'mnist_mlp_1_2048': 'VIMODEL_MNIST_1_2048_relu.net.npz',
    'mnist_mlp_1_256': 'VIMODEL_MNIST_1_256_relu.net.npz',
    'mnist_mlp_1_4096': 'VIMODEL_MNIST_1_4096_relu.net.npz',
    'mnist_mlp_1_512': 'VIMODEL_MNIST_1_512_relu.net.npz',
    'mnist_mlp_1_64': 'VIMODEL_MNIST_1_64_relu.net.npz',
    'mnist_mlp_2_1024': 'VIMODEL_MNIST_2_1024_relu.net.npz',
    'mnist_mlp_2_128': 'VIMODEL_MNIST_2_128_relu.net.npz',
    'mnist_mlp_2_256': 'VIMODEL_MNIST_2_256_relu.net.npz',
    'mnist_mlp_2_512': 'VIMODEL_MNIST_2_512_relu.net.npz',
    'mnist_mlp_2_64': 'VIMODEL_MNIST_2_64_relu.net.npz',
})


def _load_pickled_model(root_dir: str, model_name: str) -> ModelParams:
  model_path = getattr(INTERNAL_MODEL_PATHS, model_name.lower())
  if model_path.endswith('mnist_ceda.pkl'):
    with jv_utils.open_file(model_path, 'rb', root_dir=root_dir) as f:
      params_iterables = pickle.load(f, encoding='bytes')
  else:

    with jv_utils.open_file(model_path, 'rb', root_dir=root_dir) as f: