# Path: jaxopt/_src/base.py
# class OptStep(NamedTuple):
# class KKTSolution(NamedTuple):
# class Solver(abc.ABC):
# class IterativeSolver(Solver):
# class StochasticSolver(IterativeSolver):
# class LinearOperator(object):
# class LineSearchStep(NamedTuple):
# class IterativeLineSearch(IterativeSolver):
#   def run(self,
#           init_params: Any,
#           *args,
#           **kwargs) -> OptStep:
#   def l2_optimality_error(self, params, *args, **kwargs):
#   def attribute_names(self):
#   def attribute_values(self):
#   def _get_loop_options(self):
#   def _cond_fun(self, inputs):
#   def _body_fun(self, inputs):
#   def _make_zero_step(self, init_params, state) -> OptStep:
#   def _run(self,
#            init_params: Any,
#            *args,
#            **kwargs) -> OptStep:
#   def run(self,
#           init_params: Any,
#           *args,
#           **kwargs) -> OptStep:
# def _where(cond, x, y):
#   def run_iterator(self,
#                    init_params: Any,
#                    iterator,
#                    *args,
#                    **kwargs) -> OptStep:
#   def __init__(self, A):
#   def shape(self):
#   def matvec(self, x):
#   def matvec_element(self, x, idx):
#   def rmatvec(self, x):
#   def rmatvec_element(self, x, idx):
#   def update_matvec(self, Ax, delta, idx):
#   def update_rmatvec(self, ATx, delta, idx):
#   def column_l2_norms(self, squared=False):
#   def tree_flatten(self):
#   def tree_unflatten(cls, aux_data, children):
#   def _make_zero_step(self, init_stepsize, state) -> LineSearchStep:
#   def run(self,
#           init_stepsize: float,
#           params: Any,
#           value: Optional[float] = None,
#           grad: Optional[Any] = None,
#           descent_direction: Optional[Any] = None,
#           *args,
#           **kwargs) -> LineSearchStep:
# 
# Path: jaxopt/_src/implicit_diff.py
# def root_vjp(optimality_fun: Callable,
#              sol: Any,
#              args: Tuple,
#              cotangent: Any,
#              solve: Callable = linear_solve.solve_normal_cg) -> Any:
#   def fun_sol(sol):
#   def fun_args(*args):
# def _jvp_sol(optimality_fun, sol, args, tangent):
# def _jvp_args(optimality_fun, sol, args, tangents):
# def root_jvp(optimality_fun: Callable,
#              sol: Any,
#              args: Tuple,
#              tangents: Tuple,
#              solve:Callable = linear_solve.solve_normal_cg) -> Any:
# def _extract_kwargs(kwarg_keys, flat_args):
# def _signature_bind(signature, *args, **kwargs):
# def _signature_bind_and_match(signature, *args, **kwargs):
#   def map_back(out_args):
# def _custom_root(solver_fun, optimality_fun, solve, has_aux,
#                  reference_signature=None):
#   def make_custom_vjp_solver_fun(solver_fun, kwarg_keys):
#     def solver_fun_flat(*flat_args):
#     def solver_fun_fwd(*flat_args):
#     def solver_fun_bwd(tup, cotangent):
#   def wrapped_solver_fun(*args, **kwargs):
# def custom_root(optimality_fun: Callable,
#                 has_aux: bool = False,
#                 solve: Callable = linear_solve.solve_normal_cg,
#                 reference_signature: Optional[Callable] = None):
#   def wrapper(solver_fun):
# def custom_fixed_point(fixed_point_fun: Callable,
#                        has_aux: bool = False,
#                        solve: Callable = linear_solve.solve_normal_cg,
#                        reference_signature: Optional[Callable] = None):
#   def optimality_fun(params, *args):
# def make_kkt_optimality_fun(obj_fun, eq_fun, ineq_fun=None):
#   def optimality_fun(params, params_obj, params_eq, params_ineq):
# 
# Path: jaxopt/_src/loop.py
# def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
#   def _iter(val):
#   def _fun(tup, it):
# def _while_loop_python(cond_fun, body_fun, init_val, maxiter):
# def _while_loop_lax(cond_fun, body_fun, init_val, maxiter):
#   def _cond_fun(_val):
#   def _body_fun(_val):
# def while_loop(cond_fun, body_fun, init_val, maxiter, unroll=False, jit=False):


from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src import implicit_diff as idf
from jaxopt._src import loop
import jax.numpy as jnp
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://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.

"""Bisection algorithm in JAX."""






class BisectionState(NamedTuple):
  """Named tuple containing state information."""
  iter_num: int
  value: float
  error: float
  low: float
  high: float
  sign: int
  aux: Optional[Any] = None


@dataclass(eq=False)

# Path: jaxopt/_src/base.py
# class OptStep(NamedTuple):
# class KKTSolution(NamedTuple):
# class Solver(abc.ABC):
# class IterativeSolver(Solver):
# class StochasticSolver(IterativeSolver):
# class LinearOperator(object):
# class LineSearchStep(NamedTuple):
# class IterativeLineSearch(IterativeSolver):
#   def run(self,
#           init_params: Any,
#           *args,
#           **kwargs) -> OptStep:
#   def l2_optimality_error(self, params, *args, **kwargs):
#   def attribute_names(self):
#   def attribute_values(self):
#   def _get_loop_options(self):
#   def _cond_fun(self, inputs):
#   def _body_fun(self, inputs):
#   def _make_zero_step(self, init_params, state) -> OptStep:
#   def _run(self,
#            init_params: Any,
#            *args,
#            **kwargs) -> OptStep:
#   def run(self,
#           init_params: Any,
#           *args,
#           **kwargs) -> OptStep:
# def _where(cond, x, y):
#   def run_iterator(self,
#                    init_params: Any,
#                    iterator,
#                    *args,
#                    **kwargs) -> OptStep:
#   def __init__(self, A):
#   def shape(self):
#   def matvec(self, x):
#   def matvec_element(self, x, idx):
#   def rmatvec(self, x):
#   def rmatvec_element(self, x, idx):
#   def update_matvec(self, Ax, delta, idx):
#   def update_rmatvec(self, ATx, delta, idx):
#   def column_l2_norms(self, squared=False):
#   def tree_flatten(self):
#   def tree_unflatten(cls, aux_data, children):
#   def _make_zero_step(self, init_stepsize, state) -> LineSearchStep:
#   def run(self,
#           init_stepsize: float,
#           params: Any,
#           value: Optional[float] = None,
#           grad: Optional[Any] = None,
#           descent_direction: Optional[Any] = None,
#           *args,
#           **kwargs) -> LineSearchStep:
# 
# Path: jaxopt/_src/implicit_diff.py
# def root_vjp(optimality_fun: Callable,
#              sol: Any,
#              args: Tuple,
#              cotangent: Any,
#              solve: Callable = linear_solve.solve_normal_cg) -> Any:
#   def fun_sol(sol):
#   def fun_args(*args):
# def _jvp_sol(optimality_fun, sol, args, tangent):
# def _jvp_args(optimality_fun, sol, args, tangents):
# def root_jvp(optimality_fun: Callable,
#              sol: Any,
#              args: Tuple,
#              tangents: Tuple,
#              solve:Callable = linear_solve.solve_normal_cg) -> Any:
# def _extract_kwargs(kwarg_keys, flat_args):
# def _signature_bind(signature, *args, **kwargs):
# def _signature_bind_and_match(signature, *args, **kwargs):
#   def map_back(out_args):
# def _custom_root(solver_fun, optimality_fun, solve, has_aux,
#                  reference_signature=None):
#   def make_custom_vjp_solver_fun(solver_fun, kwarg_keys):
#     def solver_fun_flat(*flat_args):
#     def solver_fun_fwd(*flat_args):
#     def solver_fun_bwd(tup, cotangent):
#   def wrapped_solver_fun(*args, **kwargs):
# def custom_root(optimality_fun: Callable,
#                 has_aux: bool = False,
#                 solve: Callable = linear_solve.solve_normal_cg,
#                 reference_signature: Optional[Callable] = None):
#   def wrapper(solver_fun):
# def custom_fixed_point(fixed_point_fun: Callable,
#                        has_aux: bool = False,
#                        solve: Callable = linear_solve.solve_normal_cg,
#                        reference_signature: Optional[Callable] = None):
#   def optimality_fun(params, *args):
# def make_kkt_optimality_fun(obj_fun, eq_fun, ineq_fun=None):
#   def optimality_fun(params, params_obj, params_eq, params_ineq):
# 
# Path: jaxopt/_src/loop.py
# def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
#   def _iter(val):
#   def _fun(tup, it):
# def _while_loop_python(cond_fun, body_fun, init_val, maxiter):
# def _while_loop_lax(cond_fun, body_fun, init_val, maxiter):
#   def _cond_fun(_val):
#   def _body_fun(_val):
# def while_loop(cond_fun, body_fun, init_val, maxiter, unroll=False, jit=False):


# Path: jaxopt/_src/bisection.py
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src import implicit_diff as idf
from jaxopt._src import loop
import jax.numpy as jnp

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://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.

"""Bisection algorithm in JAX."""






class BisectionState(NamedTuple):
  """Named tuple containing state information."""
  iter_num: int
  value: float
  error: float
  low: float
  high: float
  sign: int
  aux: Optional[Any] = None


@dataclass(eq=False)

class Bisection(base.IterativeSolver):