# 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
  """Make list of LayerParams from list of iterables."""
  conv_field_names = [
      f.name for f in dataclasses.fields(verify_utils.ConvParams)
  ]
  fc_field_names = [
      f.name for f in dataclasses.fields(verify_utils.FCParams)
  ]

  net = []
  for layer_params in raw_params:
    if isinstance(layer_params, tuple):
      w, b = layer_params
      layer = verify_utils.FCParams(w=w, b=b)
    elif (isinstance(layer_params, dict)
          and layer_params.get('type') == 'linear'):
      fc_params = dict(
          (k, v) for k, v in layer_params.items() if k in fc_field_names)
      if fc_params.get('dropout_rate', 0) > 0:
        w = fc_params['w']
        # adapt expected value of 'w'
        fc_params['w'] = w * (1.0 - fc_params['dropout_rate'])
        fc_params['w_bound'] = jax_verify.IntervalBound(
            lower_bound=jnp.minimum(w, 0.0), upper_bound=jnp.maximum(w, 0.0))
      layer = verify_utils.FCParams(**fc_params)
    elif isinstance(layer_params, dict):
      conv_params = dict(
          (k, v) for k, v in layer_params.items() if k in conv_field_names)
      # deal with 'W' vs 'w'
      if 'W' in layer_params:
        conv_params['w'] = layer_params['W']

# 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

  """Make list of LayerParams from list of iterables."""
  conv_field_names = [
      f.name for f in dataclasses.fields(verify_utils.ConvParams)
  ]
  fc_field_names = [
      f.name for f in dataclasses.fields(verify_utils.FCParams)
  ]

  net = []
  for layer_params in raw_params:
    if isinstance(layer_params, tuple):
      w, b = layer_params
      layer = verify_utils.FCParams(w=w, b=b)
    elif (isinstance(layer_params, dict)
          and layer_params.get('type') == 'linear'):
      fc_params = dict(
          (k, v) for k, v in layer_params.items() if k in fc_field_names)
      if fc_params.get('dropout_rate', 0) > 0:
        w = fc_params['w']
        # adapt expected value of 'w'
        fc_params['w'] = w * (1.0 - fc_params['dropout_rate'])
        fc_params['w_bound'] = jax_verify.IntervalBound(
            lower_bound=jnp.minimum(w, 0.0), upper_bound=jnp.maximum(w, 0.0))
      layer = verify_utils.FCParams(**fc_params)
    elif isinstance(layer_params, dict):
      conv_params = dict(
          (k, v) for k, v in layer_params.items() if k in conv_field_names)
      # deal with 'W' vs 'w'
      if 'W' in layer_params:
        conv_params['w'] = layer_params['W']

      layer = verify_utils.ConvParams(**conv_params)