# 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
    return hv_flat

  def Av(v):
    """Matrix-vector product.

    Args:
      v: vector, DeviceArray

    Returns:
      Av: vector, Device array. A is defined as diag(kappa) - M(lambda) where
          M(lambda) = [0, g_lambda';
                       g_lambda, H_lambda], and these terms correspond to
          L~(z) = c_lambda + g_lambda' z + z' H_lambda z
    """
    # Expand Mv=[0 g'; g H] [v0;v1] = [g'v1; v0*g + H(v1)] = [Mv0;Mv1]
    # Compute Mv0 term
    mv_zero = jnp.reshape(jnp.vdot(g_lambda, v[1:]), (1,))
    # Compute Mv1 term
    mv_rest = Hv(v[1:]) + v[0] * g_lambda
    mv = jnp.concatenate([mv_zero, mv_rest], axis=0)
    diag_kappa_v = jnp.reshape(dual_vars[-1], mv.shape) * v
    av = diag_kappa_v - mv
    return jnp.reshape(av, v.shape)

  # Construct dual function (dual_vars[-1]=kappa)
  if exact:
    eig_vec, eig_info = eigenvector_utils.min_eigenvector_exact(
        Av, N, scl=scl, report_all=True)
    info['eig_info'] = eig_info
  else:

# 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

    return hv_flat

  def Av(v):
    """Matrix-vector product.

    Args:
      v: vector, DeviceArray

    Returns:
      Av: vector, Device array. A is defined as diag(kappa) - M(lambda) where
          M(lambda) = [0, g_lambda';
                       g_lambda, H_lambda], and these terms correspond to
          L~(z) = c_lambda + g_lambda' z + z' H_lambda z
    """
    # Expand Mv=[0 g'; g H] [v0;v1] = [g'v1; v0*g + H(v1)] = [Mv0;Mv1]
    # Compute Mv0 term
    mv_zero = jnp.reshape(jnp.vdot(g_lambda, v[1:]), (1,))
    # Compute Mv1 term
    mv_rest = Hv(v[1:]) + v[0] * g_lambda
    mv = jnp.concatenate([mv_zero, mv_rest], axis=0)
    diag_kappa_v = jnp.reshape(dual_vars[-1], mv.shape) * v
    av = diag_kappa_v - mv
    return jnp.reshape(av, v.shape)

  # Construct dual function (dual_vars[-1]=kappa)
  if exact:
    eig_vec, eig_info = eigenvector_utils.min_eigenvector_exact(
        Av, N, scl=scl, report_all=True)
    info['eig_info'] = eig_info
  else:

    eig_vec = eigenvector_utils.min_eigenvector_lanczos(