# 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
    params = make_params(next(self.prng_seq))
    self._check_deterministic_behavior(params)

  def test_forward_almost_no_randomness(self):
    params = make_params(next(self.prng_seq), std=1e-8, dropout_rate=1e-8)
    self._check_deterministic_behavior(params)

  def test_forward_gaussian(self):
    params = make_params(next(self.prng_seq), std=1.0)
    self._check_stochastic_behavior(params)

  def test_forward_dropout(self):
    params = make_params(next(self.prng_seq), dropout_rate=0.8)
    self._check_stochastic_behavior(params)

  def test_adversarial_integration(self):
    spec_type = verify_utils.SpecType.ADVERSARIAL
    params = make_params(next(self.prng_seq), std=0.1, dropout_rate=0.2)
    attacks.adversarial_attack(
        params,
        self.data_spec,
        spec_type,
        next(self.prng_seq),
        num_steps=5,
        learning_rate=0.1,
        num_samples=3)

  def test_adversarial_uncertainty_integration(self):
    spec_type = verify_utils.SpecType.ADVERSARIAL
    params = make_params(next(self.prng_seq), std=0.1, dropout_rate=0.2)

# 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

    params = make_params(next(self.prng_seq))
    self._check_deterministic_behavior(params)

  def test_forward_almost_no_randomness(self):
    params = make_params(next(self.prng_seq), std=1e-8, dropout_rate=1e-8)
    self._check_deterministic_behavior(params)

  def test_forward_gaussian(self):
    params = make_params(next(self.prng_seq), std=1.0)
    self._check_stochastic_behavior(params)

  def test_forward_dropout(self):
    params = make_params(next(self.prng_seq), dropout_rate=0.8)
    self._check_stochastic_behavior(params)

  def test_adversarial_integration(self):
    spec_type = verify_utils.SpecType.ADVERSARIAL
    params = make_params(next(self.prng_seq), std=0.1, dropout_rate=0.2)
    attacks.adversarial_attack(
        params,
        self.data_spec,
        spec_type,
        next(self.prng_seq),
        num_steps=5,
        learning_rate=0.1,
        num_samples=3)

  def test_adversarial_uncertainty_integration(self):
    spec_type = verify_utils.SpecType.ADVERSARIAL
    params = make_params(next(self.prng_seq), std=0.1, dropout_rate=0.2)

    attacks.adversarial_attack(