# Path: jax_verify/extensions/sdp_verify/eigenvector_utils.py
# def safe_eigh(a, UPLO=None, symmetrize_input=True):
# def lanczos_alg(matrix_vector_product,
#                 dim,
#                 order,
#                 rng_key,
#                 dynamic_unroll=True,
#                 use_jax=True,
#                 verbose=False):
#     def _index_update(array, index, value):
# def _lanczos_alg_dynamic_unroll(
#     matrix_vector_product, dim, order, rng_key):
#   def _body_fn_update_alpha(i, vecs, tridiag, beta):
#   def body_fn(i, vals):
# def _make_pos(vecx):
# def max_eigenvector_lanczos(matrix_vector_product, dim, order, key, scl=-1,
#                             dynamic_unroll=True, use_safe_eig_vec=True):
#     def get_eig_vec(vals):
# def min_eigenvector_lanczos(matrix_vector_product, *args, **kwargs):
# def max_eigenvector_exact(matrix_vector_product, vec_dim, scl=-1,
#                           report_all=False):
#   def batched_Hv(v_batched):
# def min_eigenvector_exact(matrix_vector_product, vec_dim, scl=-1,
#                           report_all=False):
#   H = batched_Hv(jnp.eye(vec_dim))
#   H = (H + H.T)/2
# 
# 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 collections
import functools
import jax
import jax.numpy as jnp
import jax.scipy
import numpy as np
import optax
import tree
from absl import logging
from jax_verify.extensions.sdp_verify import eigenvector_utils
from jax_verify.extensions.sdp_verify import utils

    dual_vars = project_duals(dual_vars, verif_instance.dual_types)

    if opt_dual_vars:
      distance_to_opt = jax.tree_multimap(lambda x, y: jnp.linalg.norm(x - y),
                                          dual_vars, opt_dual_vars)
      info['distance_to_opt'].append(distance_to_opt)

    if i % eval_every == 0:
      dual_val, _ = loss(dual_vars, loss_scl=-1, exact=use_exact_eig_eval)
      info['steps'].append(i)
      info['loss_vals'].append(float(dual_val))
      if verbose:
        print(f'Dual iter {i}: Train loss: {loss_val} Loss {dual_val}')

  final_loss = float(loss(dual_vars, loss_scl=-1, exact=use_exact_eig_eval)[0])
  info['final_dual_vars'] = dual_vars
  info['final_opt_state'] = opt_state
  info['final_loss'] = final_loss
  info['loss_log'] = loss_log
  info['store_best'] = store_best
  if include_opt_state:
    return final_loss, info, opt_state
  else:
    return final_loss, info

solve_dual_sdp_elided = solve_sdp_dual  # Alias

############     Dual initialization     ############


# Path: jax_verify/extensions/sdp_verify/eigenvector_utils.py
# def safe_eigh(a, UPLO=None, symmetrize_input=True):
# def lanczos_alg(matrix_vector_product,
#                 dim,
#                 order,
#                 rng_key,
#                 dynamic_unroll=True,
#                 use_jax=True,
#                 verbose=False):
#     def _index_update(array, index, value):
# def _lanczos_alg_dynamic_unroll(
#     matrix_vector_product, dim, order, rng_key):
#   def _body_fn_update_alpha(i, vecs, tridiag, beta):
#   def body_fn(i, vals):
# def _make_pos(vecx):
# def max_eigenvector_lanczos(matrix_vector_product, dim, order, key, scl=-1,
#                             dynamic_unroll=True, use_safe_eig_vec=True):
#     def get_eig_vec(vals):
# def min_eigenvector_lanczos(matrix_vector_product, *args, **kwargs):
# def max_eigenvector_exact(matrix_vector_product, vec_dim, scl=-1,
#                           report_all=False):
#   def batched_Hv(v_batched):
# def min_eigenvector_exact(matrix_vector_product, vec_dim, scl=-1,
#                           report_all=False):
#   H = batched_Hv(jnp.eye(vec_dim))
#   H = (H + H.T)/2
# 
# 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: jax_verify/extensions/sdp_verify/sdp_verify.py
import collections
import functools
import jax
import jax.numpy as jnp
import jax.scipy
import numpy as np
import optax
import tree
from absl import logging
from jax_verify.extensions.sdp_verify import eigenvector_utils
from jax_verify.extensions.sdp_verify import utils


    dual_vars = project_duals(dual_vars, verif_instance.dual_types)

    if opt_dual_vars:
      distance_to_opt = jax.tree_multimap(lambda x, y: jnp.linalg.norm(x - y),
                                          dual_vars, opt_dual_vars)
      info['distance_to_opt'].append(distance_to_opt)

    if i % eval_every == 0:
      dual_val, _ = loss(dual_vars, loss_scl=-1, exact=use_exact_eig_eval)
      info['steps'].append(i)
      info['loss_vals'].append(float(dual_val))
      if verbose:
        print(f'Dual iter {i}: Train loss: {loss_val} Loss {dual_val}')

  final_loss = float(loss(dual_vars, loss_scl=-1, exact=use_exact_eig_eval)[0])
  info['final_dual_vars'] = dual_vars
  info['final_opt_state'] = opt_state
  info['final_loss'] = final_loss
  info['loss_log'] = loss_log
  info['store_best'] = store_best
  if include_opt_state:
    return final_loss, info, opt_state
  else:
    return final_loss, info

solve_dual_sdp_elided = solve_sdp_dual  # Alias

############     Dual initialization     ############


DualVarTypes = utils.DualVarTypes