# Path: jax_verify/sdp_verify.py

# 
# Path: jax_verify/extensions/sdp_verify/boundprop_utils.py
# def boundprop(params, x, epsilon, input_bounds, boundprop_type,
#               **extra_boundprop_kwargs):
# def _crown_ibp_boundprop(params, x, epsilon, input_bounds):
#   def get_layer_act(layer_idx, inputs):
# def _nonconvex_boundprop(params, x, epsilon, input_bounds,
#                          nonconvex_boundprop_steps=100,
#                          nonconvex_boundprop_nodes=128):
# 
# Path: jax_verify/extensions/sdp_verify/problem.py
# DEFAULT_DISABLED_DUAL_VARS = ('nu_quad', 'muminus2')
# NECESSARY_DUAL_VARS = ('lam', 'muplus', 'muminus')
#   W0 = W0_orig * jnp.reshape(jnp.sqrt(sigmasq_z), (-1, 1))
#   N = sum([np.prod(np.array(i)) for i in layer_sizes])
# def make_relu_network_lagrangian(dual_vars, params, bounds, obj):
#   def lagrangian(xs_list):
# def relu_robustness_verif_instance_to_sdp(verif_instance):
#   def obj(x_final):
#   def make_inner_lagrangian(dual_vars):
# def make_sdp_verif_instance(verif_instance):
# def make_vae_sdp_verif_instance(params, data_x, bounds):
#   def recon_loss(x_final):
#   def make_inner_lagrangian(dual_vars):
# def make_vae_semantic_spec_params(x, vae_params, classifier_params):
# def get_dual_shapes_and_types(bounds_elided):
# 
# Path: jax_verify/extensions/sdp_verify/utils.py
# 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'


import functools
import os
import pickle
import jax
import jax.numpy as jnp
import jax_verify
import numpy as np
from absl import app
from absl import flags
from jax_verify import sdp_verify
from jax_verify.extensions.sdp_verify import boundprop_utils
from jax_verify.extensions.sdp_verify import problem
from jax_verify.extensions.sdp_verify import utils
      _opt_multiplier_fn, kappa_index=kappa_index, kappa_dim=kappa_shape[1])

  # Call solver
  obj_value, info = sdp_verify.solve_sdp_dual(
      verif_instance,
      num_steps=num_steps,
      verbose=True,
      opt_multiplier_fn=opt_multiplier_fn,
      **solver_params)
  info['final_dual_vars'] = jax.tree_map(np.array, info['final_dual_vars'])
  return float(obj_value), info


class PickleWriter:

  def write(self, d):
    with open('/tmp/run_sdp_verify_results.pkl', 'wb') as f:
      pickle.dump(d, f)


def main(unused_argv):
  run_verification(PickleWriter())


def run_verification(writer):
  """Run verification."""
  xs, ys = _load_dataset(FLAGS.dataset)
  dataset_idx = FLAGS.dataset_idx
  if FLAGS.dataset == 'cifar10':
    x = utils.preprocess_cifar(xs[dataset_idx])

# Path: jax_verify/sdp_verify.py

# 
# Path: jax_verify/extensions/sdp_verify/boundprop_utils.py
# def boundprop(params, x, epsilon, input_bounds, boundprop_type,
#               **extra_boundprop_kwargs):
# def _crown_ibp_boundprop(params, x, epsilon, input_bounds):
#   def get_layer_act(layer_idx, inputs):
# def _nonconvex_boundprop(params, x, epsilon, input_bounds,
#                          nonconvex_boundprop_steps=100,
#                          nonconvex_boundprop_nodes=128):
# 
# Path: jax_verify/extensions/sdp_verify/problem.py
# DEFAULT_DISABLED_DUAL_VARS = ('nu_quad', 'muminus2')
# NECESSARY_DUAL_VARS = ('lam', 'muplus', 'muminus')
#   W0 = W0_orig * jnp.reshape(jnp.sqrt(sigmasq_z), (-1, 1))
#   N = sum([np.prod(np.array(i)) for i in layer_sizes])
# def make_relu_network_lagrangian(dual_vars, params, bounds, obj):
#   def lagrangian(xs_list):
# def relu_robustness_verif_instance_to_sdp(verif_instance):
#   def obj(x_final):
#   def make_inner_lagrangian(dual_vars):
# def make_sdp_verif_instance(verif_instance):
# def make_vae_sdp_verif_instance(params, data_x, bounds):
#   def recon_loss(x_final):
#   def make_inner_lagrangian(dual_vars):
# def make_vae_semantic_spec_params(x, vae_params, classifier_params):
# def get_dual_shapes_and_types(bounds_elided):
# 
# Path: jax_verify/extensions/sdp_verify/utils.py
# 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'


# Path: examples/run_sdp_verify.py
import functools
import os
import pickle
import jax
import jax.numpy as jnp
import jax_verify
import numpy as np
from absl import app
from absl import flags
from jax_verify import sdp_verify
from jax_verify.extensions.sdp_verify import boundprop_utils
from jax_verify.extensions.sdp_verify import problem
from jax_verify.extensions.sdp_verify import utils

      _opt_multiplier_fn, kappa_index=kappa_index, kappa_dim=kappa_shape[1])

  # Call solver
  obj_value, info = sdp_verify.solve_sdp_dual(
      verif_instance,
      num_steps=num_steps,
      verbose=True,
      opt_multiplier_fn=opt_multiplier_fn,
      **solver_params)
  info['final_dual_vars'] = jax.tree_map(np.array, info['final_dual_vars'])
  return float(obj_value), info


class PickleWriter:

  def write(self, d):
    with open('/tmp/run_sdp_verify_results.pkl', 'wb') as f:
      pickle.dump(d, f)


def main(unused_argv):
  run_verification(PickleWriter())


def run_verification(writer):
  """Run verification."""
  xs, ys = _load_dataset(FLAGS.dataset)
  dataset_idx = FLAGS.dataset_idx
  if FLAGS.dataset == 'cifar10':
    x = utils.preprocess_cifar(xs[dataset_idx])

    epsilon, input_bounds = utils.preprocessed_cifar_eps_and_input_bounds(