class DualOp(bound_propagation.Bound):
class _LagrangianTransform(bound_propagation.GraphTransform[DualOp]):
class InnerMaxStrategy(metaclass=abc.ABCMeta):
  def __init__(
      self,
      name,
      base_bound: bound_propagation.Bound,
      affine_fn: Callable[[jnp.array], jnp.array],
      inputs: Optional[Sequence[Union['DualOp', jnp.array]]] = None,
      relu_preact_name: Optional[int] = None):
  def base_bound(self) -> bound_propagation.Bound:
  def lower(self) -> jnp.array:
  def upper(self) -> jnp.array:
  def shape(self) -> Sequence[int]:
  def affine(self, act_or_input):
  def is_input(self) -> bool:
  def is_relu(self) -> bool:
  def relu_preact_name(self) -> int:
  def inputs(self) -> Sequence[Union['DualOp', jnp.array]]:
  def __init__(self, boundprop_transform: bound_propagation.BoundTransform):
  def input_transform(self, context, lower_bound, upper_bound):
  def primitive_transform(self, context, primitive, *args, **params):
      def affine_fn(act_or_input):
  def solve_max(
      self,
      inner_dual_vars: Any,
      opt_instance: InnerVerifInstance,
      key: jnp.array,
      step: int,
  ) -> jnp.array:
  def supports_stochastic_parameters(self):
  def build_spec(self, opt_instance: InnerVerifInstance, step: int,
                 softmax: bool = False):
    def forward_relu_before_affine(x):
    def forward_affine_before_relu(x):
    def obj_first(x, duals_pre, duals_post):
    def obj_intermediate(x, duals_pre, duals_post):
    def obj_last(x, duals_pre, duals_post):
  def init_duals(
      self,
      boundprop_transform: bound_propagation.BoundTransform,
      spec_type: verify_utils.SpecType,
      affine_before_relu: bool,
      spec_fn: Callable[..., jnp.array],
      key: jnp.array,
      lagrangian_form_per_layer: Iterable[LagrangianForm],
      *input_bounds: Nest[graph_traversal.GraphInput],
  ) -> Tuple[Dict[int, DualOp], Params, ParamsTypes]:
  def init_layer_inner_params(
      self, opt_instance: verify_utils.InnerVerifInstance) -> Tuple[Any, Any]:
def project_dual(dual_params: Params,
                 dual_params_types: ParamsTypes) -> Params:
def build_dual_fun(
    env: Dict[int, DualOp],
    lagrangian_form: Dict[int, LagrangianForm],
    inner_opt: InnerMaxStrategy,
    affine_before_relu: bool,
    spec_type: verify_utils.SpecType,
    merge_problems: Optional[Dict[int, int]] = None,
) -> Callable[[Params, jnp.array, int], jnp.array]:
  def dual_loss_fun(dual_params: Params,
                    key: jnp.array, step: int) -> jnp.array:
def _enumerate_inner_max_problems(
    env: Dict[int, DualOp],
    lagrangian_form: Dict[int, LagrangianForm],
    lagrange_params: Dict[int, Any],
    spec_type: verify_utils.SpecType,
    affine_before_relu: bool,
) -> List[InnerVerifInstance]:
def _merge_specified_instances(
    instances: Sequence[InnerVerifInstance],
    merge_specification: Dict[int, int],
) -> Sequence[InnerVerifInstance]:
def _merge_instances(
    instance_first: InnerVerifInstance,
    *instances_rest: InnerVerifInstance,
) -> InnerVerifInstance:
def make_opt_and_num_steps(opt_config):
  def lr_schedule(t): def solve_dual(
    config: ConfigDict,
    bounds: Sequence[sdp_utils.IntBound],
    spec_type: verify_utils.SpecType,
    spec_fn: Callable[..., jnp.array],
    params: ModelParams,
    dual_state: ConfigDict,
    mode: str,
    logger: Callable[[int, Mapping[str, Any]], None],
) -> Tuple[float, Tensor]:
def solve_dual_train(
    env: Dict[int, DualOp],
    dual_state: ConfigDict,
    opt: optax.GradientTransformation,
    inner_opt: InnerMaxStrategy,
    dual_params: Params,
    spec_type: verify_utils.SpecType,
    dual_params_types: ParamsTypes,
    logger: Callable[[int, Mapping[str, Any]], None],
    key: jnp.array,
    num_steps: int,
    affine_before_relu: bool,
    device_type=None,
    merge_problems: Optional[Dict[int, int]] = None,
    block_to_time: bool = False,
) -> ConfigDict:
  def grad_step(params, opt_state, key, step):
def solve_dual_eval(
    env: Dict[int, DualOp],
    inner_opt: InnerMaxStrategy,
    dual_params: Params,
    spec_type: verify_utils.SpecType,
    dual_params_types: ParamsTypes,
    logger: Callable[[int, Mapping[str, Any]], None],
    key: jnp.array,
    affine_before_relu: bool,
    step: int,
    merge_problems: Optional[Dict[int, int]] = None,
) -> float: def _flatten_spatial_dims(x: Tensor) -> Tensor:
def size_from_shape(shape: Shape) -> int:
  def __init__(self, name):
  def _init_params_per_sample(self, key: PRNGKey, *args) -> Params:
  def init_params(self, key, *args, **kwargs):
  def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
  def apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
  def process_params(self, lagrange_params: Params):
  def name(self):
  def __init__(self):
  def _init_params_per_sample(self,
                              key: PRNGKey,
                              l_shape: Shape,
                              init_zeros: bool = True) -> Params:
  def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
                        step: int) -> Tensor:
  def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
  def __init__(self):
  def _init_params_per_sample(self,
                              key: PRNGKey,
                              l_shape: Shape,
                              init_zeros: bool = False) -> Params:
  def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
                        step: int) -> Tensor:
  def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
def get_lagrangian_form(config_lagrangian_form: ConfigDict) -> LagrangianForm:
class LagrangianForm(metaclass=abc.ABCMeta):
class Linear(LagrangianForm):
class LinearExp(LagrangianForm): 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' class LpStrategy(dual_build.InnerMaxStrategy):
  def supports_stochastic_parameters(self):
  def solve_max(
      self,
      inner_dual_vars: Any,
      opt_instance: InnerVerifInstance,
      key: jnp.array,
      step: int,
  ) -> jnp.array:
  def init_layer_inner_params(self, opt_instance): class SdpDualVerifInstance(_SdpDualVerifInstance):
class VerifInstanceTypes(enum.Enum):
class DualVarTypes(enum.Enum):
def mlp_layer_sizes(params):
def nn_layer_sizes(params):
def layer_sizes_from_bounds(bounds):
def predict_mlp(params, inputs):
def fwd(inputs, layer_params):
def predict_cnn(params, inputs, include_preactivations=False):
def get_network_activs(params, x):
def get_layer_params(fun_to_extract, example_input):
  def _get_const_input_arg(eqn):
def boundprop(params, bounds_in):
def init_bound(x, epsilon, input_bounds=(0., 1.), add_batch_dim=True):
def ibp_bound_elided(verif_instance):
def ibp_bound_nonelided(verif_instance):
def make_relu_robust_verif_instance(
    params, bounds=None, target_label=1, label=2, input_bounds=None):
def make_relu_robust_verif_instance_elided(
    params, bounds=None, input_bounds=None):
def output_size_and_verif_type(params, bounds):
def elide_params(params, label, target_label, op_size):
def preprocess_cifar(image, inception_preprocess=False, perturbation=False):
def preprocessed_cifar_eps_and_input_bounds(
    shape=(32, 32, 3), epsilon=2/255, inception_preprocess=False):
def adv_objective(model_fn, x, label, target_label):
def fgsm_single(model_fn, x, label, target_label, epsilon, num_steps,
                step_size, input_bounds=(0., 1.)):
def untargeted_margin_loss(logits, labels):
def pgd_default(model_fn, x, label, epsilon, num_steps, step_size,
                input_bounds=(0., 1.)):
def pgd(adv_loss, x_init, epsilon, num_steps, step_size, input_bounds=(0., 1.)):
def scale_by_variable_opt(multipliers):
  def init_fn(params):
  def update_fn(updates, _, params=None):
def flatten(pytree, backend=np):
def unflatten_like(a, pytree):
def structure_like(tree1, tree2):
  MLP_ELIDED = 'mlp_elided'
  CNN_ELIDED = 'cnn_elided'
  EQUALITY = 'equality'
  INEQUALITY = 'inequality' def bound_propagation(
    prop_alg: PropagationAlgorithm[Repr],
    function: Callable[..., Nest[Tensor]],
    *bounds: Nest[GraphInput],
    graph_simplifier=synthetic_primitives.default_simplifier,
) -> Tuple[
    Nest[Union[Repr, Tensor]],
    Dict[jax.core.Var, Union[Repr, Tensor, Bound]]]:
  """Performs Bound Propagation on the model implemented by `function`.

  Args:
    prop_alg: Algorithm specifying how to traverse the graph and how to
      transform each node.
    function: Pure function inputs -> outputs. If the function to propagate
      through has a more complex signature, the use of `functools.partial` can
      solve that problem.
    *bounds: Nest of `IntervalBound` objects containing the lower and upper
      bounds on all the inputs, or `Tensor`s containing known inputs directly.
    graph_simplifier: Function transforming the JaxPR graph into a simpler
      graph. Default value is a function identifying specific activation
      functions, followed by grouping of linear sequences and quadratic forms.
  Returns:
    bounds: Bounds over all the outputs of the function, with the same structure
      as the output of `function`
    env: Mapping from the node of the computations to their representation.
  """
  # Replace all the jittable bounds by standard bound object.
  bounds = unjit_inputs(*bounds)

  # Parse the computation graph.
  placeholder_inputs = jax.tree_util.tree_map(
      lambda b: b.lower if isinstance(b, Bound) else b,
      bounds)
  jaxpr_maker = jax.make_jaxpr(function)
  parsed = jaxpr_maker(*placeholder_inputs)
  output_shapes = jax.eval_shape(function, *placeholder_inputs)

  flat_is_bound, _ = jax.tree_util.tree_flatten(
      jax.tree_util.tree_map(lambda b: isinstance(b, Bound), bounds))
  inp_is_bound = {var: is_bound
                  for var, is_bound in zip(parsed.jaxpr.invars, flat_is_bound)}
  simplified_graph = synthetic_primitives.simplify_graph(
      graph_simplifier, parsed.jaxpr, inp_is_bound)
  graph = PropagationGraph(simplified_graph, parsed.literals)

  outvals, env = prop_alg.propagate(graph, bounds)

  # Make outvals into the same tree structure than the output of the function.
  tree_structure = jax.tree_util.tree_structure(output_shapes)
  outvals = jax.tree_util.tree_unflatten(tree_structure, outvals)

  return outvals, env _EPS = 1e-10
class CvxpySolver(relaxation.RelaxationSolver):
class CvxpyMIPSolver(CvxpySolver, relaxation.MIPSolver):
  def __init__(self):
  def _variable_already_created(self, var: Variable) -> bool:
  def _create_solver_relax_variable(
      self,
      relax_var: relaxation.RelaxVariable,
      index: int):
  def create_linear_solver_constraint(self,
                                      constraint: relaxation.LinearConstraint,
                                      index: int):
  def create_activation_solver_constraint(
      self, constraint: relaxation.RelaxActivationConstraint, act_index: int,
      slope: float, bias: float):
  def minimize_objective(
      self,
      var_name: str,
      objective: Tensor,
      objective_bias: float,
      time_limit: Optional[int] = None,
  ) -> Tuple[float, Solution, bool]:
  def __init__(self):
  def _create_solver_bool_variable(
      self,
      bin_var: relaxation.BinaryVariable,
      index: int):
  def create_mip_activation_solver_constraint(
      self,
      constraint: relaxation.MIPActivationConstraint,
      act_index: int,
      binslope: float,
      slope: float,
      bias: float,
    ): class RelaxVariable(bound_propagation.Bound):
class OptRelaxVariable(RelaxVariable):
class BinaryVariable():
class MIPActivationConstraint:
class LinearConstraint:
class RelaxActivationConstraint:
class RelaxationSolver(metaclass=abc.ABCMeta):
class MIPSolver(RelaxationSolver):
class RelaxationTransform(bound_propagation.GraphTransform[RelaxVariable]):
class OptimizedRelaxationTransform(
    bound_propagation.GraphTransform[OptRelaxVariable]):
  def __init__(self, idx: bound_propagation.Index, base_bound):
  def set_constraints(self, constraints):
  def lower(self):
  def upper(self):
  def __init__(self, base_relax_variable, optimize_transform):
  def shape(self):
  def lower(self):
  def upper(self):
  def _optimized_bounds(self):
  def __init__(self, idx, shape):
  def __init__(self, outvar, invar, binvar, mask, binscale, scale, bias, sense):
  def encode_into_solver(self, solver: 'MIPSolver', index: int):
  def __init__(self, vars_and_coeffs, bias, sense):
  def bias(self, index: int):
  def vars_and_coeffs(self, index: int):
  def encode_into_solver(self, solver: 'RelaxationSolver', index: int):
  def __init__(self, outvar, invar, mask, scale, bias, sense):
  def encode_into_solver(self, solver: 'RelaxationSolver', index: int):
  def maybe_create_solver_variable(
      self,
      var: RelaxVariable,
      index: int):
  def _create_solver_relax_variable(
      self,
      var: RelaxVariable,
      index: int):
  def _variable_already_created(
      self,
      var: Union[RelaxVariable, BinaryVariable],
      ) -> bool:
  def create_linear_solver_constraint(
      self,
      constraint: LinearConstraint,
      index: int):
  def create_activation_solver_constraint(
      self,
      constraint: RelaxActivationConstraint,
      act_index: int,
      slope: float,
      bias: float):
  def minimize_objective(
      self,
      var_name: str,
      objective: Tensor,
      objective_bias: float,
      time_limit_millis: Optional[int],
  ) -> Tuple[float, Dict[str, Tensor], bool]:
  def maybe_create_solver_variable(
      self,
      var: Union[RelaxVariable, BinaryVariable],
      index: int):
  def _create_solver_bool_variable(
      self,
      var: BinaryVariable,
      index: int):
  def create_mip_activation_solver_constraint(
      self,
      constraint: MIPActivationConstraint,
      act_index: int,
      binslope: float,
      slope: float,
      bias: float):
def encode_relaxation(
    solver_ctor: Callable[[], RelaxationSolver],
    env: Dict[jax.core.Var, Union[RelaxVariable, Tensor]],
    index: int,
) -> RelaxationSolver:
def solve_relaxation(
    solver_ctor: Callable[[], RelaxationSolver],
    objective: Tensor,
    objective_bias: float,
    variable_opt: RelaxVariable,
    env: Dict[jax.core.Var, Union[RelaxVariable, Tensor]],
    index: int,
    time_limit_millis: Optional[int] = None,
) -> Tuple[float, Dict[str, Tensor], bool]:
def _get_linear(primitive, outval, *eqn_invars, **params):
  def funx(x):
  def fungrad(i, args):
def _get_relu_relax(lower, upper):
def _relax_input(
    index: bound_propagation.Index, in_bounds: bound_propagation.Bound,
    ) -> RelaxVariable:
def _relax_primitive(
    index: bound_propagation.Index, out_bounds: bound_propagation.Bound,
    primitive: jax.core.Primitive,
    *args, use_mip: bool = False, **kwargs
    ) -> RelaxVariable:
  def __init__(
      self,
      boundprop_transform: bound_propagation.BoundTransform,
      use_mip: bool = False,
      ):
  def input_transform(self, context, lower_bound, upper_bound):
  def primitive_transform(self, context, primitive, *args, **params):
  def __init__(
      self,
      transform: bound_propagation.GraphTransform[RelaxVariable],
      solver_ctor: Callable[[], RelaxationSolver],
      time_limit_millis: Optional[int] = None):
  def tight_bounds(self, variable: RelaxVariable) -> ibp.IntervalBound:
  def input_transform(
      self,
      context: bound_propagation.TransformContext,
      lower_bound: Tensor,
      upper_bound: Tensor,
  ) -> RelaxVariable:
  def primitive_transform(
      self,
      context: bound_propagation.TransformContext,
      primitive: jax.core.Primitive,
      *args: Union[RelaxVariable, Tensor],
      **params,
  ) -> RelaxVariable: def _init_fc_layer(key, n_in, n_out):
def _init_conv_layer(key, n_h, n_w, n_cout, n_cin):
def make_mlp_params(layer_sizes, key):
def make_cnn_params(layer_sizes, key):
def make_toy_verif_instance(seed=None, label=None, target_label=None, nn='mlp'):
def make_mlp_layer_from_conv_layer(layer_params, input_bounds):
def make_mlp_verif_instance_from_cnn(verif_instance):
  W = random.normal(k1, (n_in, n_out))/ n_in
  W = random.normal(k1, (n_h, n_w, n_cin, n_cout))
from absl.testing import absltest
from jax_verify.extensions.functional_lagrangian import dual_build
from jax_verify.extensions.functional_lagrangian import dual_solve
from jax_verify.extensions.functional_lagrangian import lagrangian_form
from jax_verify.extensions.functional_lagrangian import verify_utils
from jax_verify.extensions.functional_lagrangian.inner_solvers import lp
from jax_verify.extensions.sdp_verify import utils as sdp_utils
from jax_verify.src import bound_propagation
from jax_verify.src.mip_solver import cvxpy_relaxation_solver
from jax_verify.src.mip_solver import relaxation
from jax_verify.tests.sdp_verify import test_utils as sdp_test_utils
import chex
import jax
import jax.numpy as jnp
import jax_verify
import ml_collections
import numpy as np

NUM_SAMPLES = 1
LAYER_SIZES = [3, 4, 5, 6]


def create_inputs(prng_key):
  return jax.random.uniform(
      prng_key, [NUM_SAMPLES, LAYER_SIZES[0]], minval=0.0, maxval=1.0)


def make_model_fn(params):

  def model_fn(inputs):
    inputs = np.reshape(inputs, (inputs.shape[0], -1))
    return sdp_utils.predict_mlp(params, inputs)

  return model_fn


def get_config():
  config = ml_collections.ConfigDict()

  config.outer_opt = ml_collections.ConfigDict()
  config.outer_opt.lr_init = 0.001
  config.outer_opt.steps_per_anneal = 500
  config.outer_opt.anneal_lengths = ''
  config.outer_opt.anneal_factor = 0.1
  config.outer_opt.num_anneals = 1
  config.outer_opt.opt_name = 'adam'
  config.outer_opt.opt_kwargs = {}

  return config


class LinearTest(chex.TestCase):

  def setUp(self):
    super(LinearTest, self).setUp()

    self.target_label = 1
    self.label = 0
    self.input_bounds = (0.0, 1.0)
    self.layer_sizes = LAYER_SIZES
    self.eps = 0.1

    prng_key = jax.random.PRNGKey(13579)

    self.keys = jax.random.split(prng_key, 5)
    self.network_params = sdp_test_utils.make_mlp_params(
        self.layer_sizes, self.keys[0])

    self.inputs = create_inputs(self.keys[1])

    objective = jnp.zeros(self.layer_sizes[-1])
    objective = objective.at[self.target_label].add(1)
    objective = objective.at[self.label].add(-1)
    self.objective = objective
    self.objective_bias = jax.random.normal(self.keys[2], [])

  def solve_with_jax_verify(self):
    lower_bound = jnp.minimum(jnp.maximum(self.inputs - self.eps, 0.0), 1.0)
    upper_bound = jnp.minimum(jnp.maximum(self.inputs + self.eps, 0.0), 1.0)
    init_bound = jax_verify.IntervalBound(lower_bound, upper_bound)

    logits_fn = make_model_fn(self.network_params)

    solver = cvxpy_relaxation_solver.CvxpySolver
    relaxation_transform = relaxation.RelaxationTransform(
        jax_verify.ibp_transform)

    var, env = bound_propagation.bound_propagation(
        bound_propagation.ForwardPropagationAlgorithm(relaxation_transform),
        logits_fn, init_bound)

    # This solver minimizes the objective -> get max with -min(-objective)
    neg_value_opt, _, _ = relaxation.solve_relaxation(
        solver,
        -self.objective,
        -self.objective_bias,
        var,
        env,
        index=0,
        time_limit_millis=None)
    value_opt = -neg_value_opt

    return value_opt

  def solve_with_functional_lagrangian(self):
    config = get_config()

    init_bound = sdp_utils.init_bound(
        self.inputs[0], self.eps, input_bounds=self.input_bounds)
    bounds = sdp_utils.boundprop(
        self.network_params + [(self.objective, self.objective_bias)],
        init_bound)

    logits_fn = make_model_fn(self.network_params)

    def spec_fn(inputs):
      return jnp.matmul(logits_fn(inputs), self.objective) + self.objective_bias

    input_bounds = jax_verify.IntervalBound(bounds[0].lb, bounds[0].ub)

    lagrangian_form_per_layer = lagrangian_form.Linear()
    lagrangian_form_per_layer = [lagrangian_form_per_layer for bd in bounds]
    inner_opt = lp.LpStrategy()
    env, dual_params, dual_params_types = inner_opt.init_duals(
        jax_verify.ibp_transform, verify_utils.SpecType.ADVERSARIAL, False,
        spec_fn, self.keys[3], lagrangian_form_per_layer, input_bounds)
    opt, num_steps = dual_build.make_opt_and_num_steps(config.outer_opt)
    dual_state = ml_collections.ConfigDict(type_safe=False)
    dual_solve.solve_dual_train(
        env,
        key=self.keys[4],
        num_steps=num_steps,
        opt=opt,
        dual_params=dual_params,
        dual_params_types=dual_params_types,
        dual_state=dual_state,
        affine_before_relu=False,

        spec_type=verify_utils.SpecType.ADVERSARIAL,