# 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
# coding=utf-8
# Copyright 2021 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Attacks test."""


EPS = 0.1


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


# 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

# coding=utf-8
# Copyright 2021 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Attacks test."""


EPS = 0.1


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


  return verify_utils.DataSpec(