# Path: jax_verify/extensions/functional_lagrangian/attacks.py
# def sample_truncated_normal(
#     mean: Tensor,
#     std: Tensor,
#     bounds: IntervalBound,
#     prng_key: PRNGKey
# ) -> Tensor:
# def sample_dropout(
#     mean: Tensor,
#     dropout_rate: float,
#     prng_key: PRNGKey,
# ) -> Tensor:
# def make_params_sampling_fn(
#     params: Union[ModelParams, ModelParamsElided],
# ) -> Callable[[PRNGKey], Union[ModelParams, ModelParamsElided]]:
#   def sample_fn(key: PRNGKey):
# def make_forward(
#     model_params: ModelParams,
#     num_samples: int,
# ) -> Callable[[Tensor, PRNGKey], Tensor]:
#   def single_forward(inputs, prng_key):
#   def multiple_forward(inputs, prng_key):
# def _run_attack(
#     max_objective_fn: Callable[[Tensor, PRNGKey], Tensor],
#     projection_fn: Callable[[Tensor], Tensor],
#     x_init: Tensor,
#     prng_key: PRNGKey,
#     num_steps: int,
#     learning_rate: float,
# ):
#   def body_fn(it, inputs):
# def adversarial_attack(
#     params: ModelParams,
#     data_spec: DataSpec,
#     spec_type: verify_utils.SpecType,
#     key: PRNGKey,
#     num_steps: int,
#     learning_rate: float,
#     num_samples: int = 1,
# ) -> float:
#   def max_objective_fn_uncertainty(x, prng_key):
#   def max_objective_fn_adversarial(x, prng_key):
#   def max_objective_fn_adversarial_softmax(x, prng_key):
# 
# 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'


from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.extensions.functional_lagrangian import attacks
from jax_verify.extensions.functional_lagrangian import verify_utils
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import jax_verify
  """Create data specification from config."""
  x = jax.random.normal(prng_key, [8])
  input_bounds = (x - EPS, x + EPS)

  return verify_utils.DataSpec(
      input=x,
      true_label=0,
      target_label=1,
      epsilon=EPS,
      input_bounds=input_bounds)


def make_params(prng_key, dropout_rate=0.0, std=None):
  prng_key_seq = hk.PRNGSequence(prng_key)

  w1 = jax.random.normal(next(prng_key_seq), [8, 4])
  b1 = jax.random.normal(next(prng_key_seq), [4])

  w2 = jax.random.normal(next(prng_key_seq), [4, 2])
  b2 = jax.random.normal(next(prng_key_seq), [2])

  if std is not None:
    w1_std = std * jnp.ones([8, 4])
    b1_std = std * jnp.ones([4])
    w1_bound = jax_verify.IntervalBound(w1 - 3 * w1_std, w1 + 3 * w1_std)
    b1_bound = jax_verify.IntervalBound(b1 - 3 * b1_std, b1 + 3 * b1_std)
  else:
    w1_std, b1_std, w1_bound, b1_bound = None, None, None, None

  params = [

# Path: jax_verify/extensions/functional_lagrangian/attacks.py
# def sample_truncated_normal(
#     mean: Tensor,
#     std: Tensor,
#     bounds: IntervalBound,
#     prng_key: PRNGKey
# ) -> Tensor:
# def sample_dropout(
#     mean: Tensor,
#     dropout_rate: float,
#     prng_key: PRNGKey,
# ) -> Tensor:
# def make_params_sampling_fn(
#     params: Union[ModelParams, ModelParamsElided],
# ) -> Callable[[PRNGKey], Union[ModelParams, ModelParamsElided]]:
#   def sample_fn(key: PRNGKey):
# def make_forward(
#     model_params: ModelParams,
#     num_samples: int,
# ) -> Callable[[Tensor, PRNGKey], Tensor]:
#   def single_forward(inputs, prng_key):
#   def multiple_forward(inputs, prng_key):
# def _run_attack(
#     max_objective_fn: Callable[[Tensor, PRNGKey], Tensor],
#     projection_fn: Callable[[Tensor], Tensor],
#     x_init: Tensor,
#     prng_key: PRNGKey,
#     num_steps: int,
#     learning_rate: float,
# ):
#   def body_fn(it, inputs):
# def adversarial_attack(
#     params: ModelParams,
#     data_spec: DataSpec,
#     spec_type: verify_utils.SpecType,
#     key: PRNGKey,
#     num_steps: int,
#     learning_rate: float,
#     num_samples: int = 1,
# ) -> float:
#   def max_objective_fn_uncertainty(x, prng_key):
#   def max_objective_fn_adversarial(x, prng_key):
#   def max_objective_fn_adversarial_softmax(x, prng_key):
# 
# 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/tests/functional_lagrangian/attacks_test.py
from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.extensions.functional_lagrangian import attacks
from jax_verify.extensions.functional_lagrangian import verify_utils
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import jax_verify

  """Create data specification from config."""
  x = jax.random.normal(prng_key, [8])
  input_bounds = (x - EPS, x + EPS)

  return verify_utils.DataSpec(
      input=x,
      true_label=0,
      target_label=1,
      epsilon=EPS,
      input_bounds=input_bounds)


def make_params(prng_key, dropout_rate=0.0, std=None):
  prng_key_seq = hk.PRNGSequence(prng_key)

  w1 = jax.random.normal(next(prng_key_seq), [8, 4])
  b1 = jax.random.normal(next(prng_key_seq), [4])

  w2 = jax.random.normal(next(prng_key_seq), [4, 2])
  b2 = jax.random.normal(next(prng_key_seq), [2])

  if std is not None:
    w1_std = std * jnp.ones([8, 4])
    b1_std = std * jnp.ones([4])
    w1_bound = jax_verify.IntervalBound(w1 - 3 * w1_std, w1 + 3 * w1_std)
    b1_bound = jax_verify.IntervalBound(b1 - 3 * b1_std, b1 + 3 * b1_std)
  else:
    w1_std, b1_std, w1_bound, b1_bound = None, None, None, None

  params = [

      verify_utils.FCParams(