# Path: jax_verify/src/mip_solver/relaxation.py
# 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:


from typing import Dict, List, Optional, Tuple, Union
from absl import logging
from cvxpy.reductions.solvers.defines import INSTALLED_MI_SOLVERS as MIP_SOLVERS
from jax_verify.src.mip_solver import relaxation
import cvxpy as cp
import jax
import jax.numpy as jnp
import numpy as np
# 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.

# Lint as: python3
"""Routines to solve the subproblem during verification."""


_EPS = 1e-10
CvxpyConstraint = cp.constraints.constraint.Constraint
Tensor = np.ndarray


# Path: jax_verify/src/mip_solver/relaxation.py
# 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:


# Path: jax_verify/src/mip_solver/cvxpy_relaxation_solver.py
from typing import Dict, List, Optional, Tuple, Union
from absl import logging
from cvxpy.reductions.solvers.defines import INSTALLED_MI_SOLVERS as MIP_SOLVERS
from jax_verify.src.mip_solver import relaxation
import cvxpy as cp
import jax
import jax.numpy as jnp
import numpy as np

# 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.

# Lint as: python3
"""Routines to solve the subproblem during verification."""


_EPS = 1e-10
CvxpyConstraint = cp.constraints.constraint.Constraint
Tensor = np.ndarray


Variable = Union[relaxation.RelaxVariable,