# 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'


import dataclasses
import jax.numpy as jnp
import jax_verify
import ml_collections
import numpy as np
from jax_verify.extensions.functional_lagrangian import verify_utils
  last_params_elided = verify_utils.FCParams(w_elided, b_elided)

  if last_params.has_bounds:
    w_bound_elided, b_bound_elided = jax_verify.interval_bound_propagation(
        elide_fn, last_params.w_bound, last_params.b_bound)
    last_params_elided = dataclasses.replace(
        last_params_elided, w_bound=w_bound_elided, b_bound=b_bound_elided)

  params_elided = params[:-1] + [last_params_elided]
  return params_elided


def elide_adversarial_softmax_spec(
    params: ModelParams,
    data_spec: DataSpec,
) -> ModelParamsElided:
  """Elide params to have uncertainty objective appended as a new last layer.

  Args:
    params: parameters of the model under verification.
    data_spec: data specification.

  Returns:
    params_elided: parameters with the uncertainty objective appended as
      the last 'layer'.
  """
  op_size = params[-1].w.shape[-1]
  e = np.zeros((op_size, 1))
  e[data_spec.target_label] = 1.
  e[data_spec.true_label] = -1.

# 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/functional_lagrangian/specification.py
import dataclasses
import jax.numpy as jnp
import jax_verify
import ml_collections
import numpy as np
from jax_verify.extensions.functional_lagrangian import verify_utils

  last_params_elided = verify_utils.FCParams(w_elided, b_elided)

  if last_params.has_bounds:
    w_bound_elided, b_bound_elided = jax_verify.interval_bound_propagation(
        elide_fn, last_params.w_bound, last_params.b_bound)
    last_params_elided = dataclasses.replace(
        last_params_elided, w_bound=w_bound_elided, b_bound=b_bound_elided)

  params_elided = params[:-1] + [last_params_elided]
  return params_elided


def elide_adversarial_softmax_spec(
    params: ModelParams,
    data_spec: DataSpec,
) -> ModelParamsElided:
  """Elide params to have uncertainty objective appended as a new last layer.

  Args:
    params: parameters of the model under verification.
    data_spec: data specification.

  Returns:
    params_elided: parameters with the uncertainty objective appended as
      the last 'layer'.
  """
  op_size = params[-1].w.shape[-1]
  e = np.zeros((op_size, 1))
  e[data_spec.target_label] = 1.
  e[data_spec.true_label] = -1.

  params_elided = params + [verify_utils.FCParams(jnp.array(e), jnp.zeros(()))]