# Path: jax_verify/src/bound_propagation.py
# def bound_propagation(
#     prop_alg: PropagationAlgorithm[Repr],
#     function: Callable[..., Nest[Tensor]],
#     *bounds: Nest[GraphInput],
#     graph_simplifier=synthetic_primitives.default_simplifier,
# ) -> Tuple[
#     Nest[Union[Repr, Tensor]],
#     Dict[jax.core.Var, Union[Repr, Tensor, Bound]]]:
#   """Performs Bound Propagation on the model implemented by `function`.
# 
#   Args:
#     prop_alg: Algorithm specifying how to traverse the graph and how to
#       transform each node.
#     function: Pure function inputs -> outputs. If the function to propagate
#       through has a more complex signature, the use of `functools.partial` can
#       solve that problem.
#     *bounds: Nest of `IntervalBound` objects containing the lower and upper
#       bounds on all the inputs, or `Tensor`s containing known inputs directly.
#     graph_simplifier: Function transforming the JaxPR graph into a simpler
#       graph. Default value is a function identifying specific activation
#       functions, followed by grouping of linear sequences and quadratic forms.
#   Returns:
#     bounds: Bounds over all the outputs of the function, with the same structure
#       as the output of `function`
#     env: Mapping from the node of the computations to their representation.
#   """
#   # Replace all the jittable bounds by standard bound object.
#   bounds = unjit_inputs(*bounds)
# 
#   # Parse the computation graph.
#   placeholder_inputs = jax.tree_util.tree_map(
#       lambda b: b.lower if isinstance(b, Bound) else b,
#       bounds)
#   jaxpr_maker = jax.make_jaxpr(function)
#   parsed = jaxpr_maker(*placeholder_inputs)
#   output_shapes = jax.eval_shape(function, *placeholder_inputs)
# 
#   flat_is_bound, _ = jax.tree_util.tree_flatten(
#       jax.tree_util.tree_map(lambda b: isinstance(b, Bound), bounds))
#   inp_is_bound = {var: is_bound
#                   for var, is_bound in zip(parsed.jaxpr.invars, flat_is_bound)}
#   simplified_graph = synthetic_primitives.simplify_graph(
#       graph_simplifier, parsed.jaxpr, inp_is_bound)
#   graph = PropagationGraph(simplified_graph, parsed.literals)
# 
#   outvals, env = prop_alg.propagate(graph, bounds)
# 
#   # Make outvals into the same tree structure than the output of the function.
#   tree_structure = jax.tree_util.tree_structure(output_shapes)
#   outvals = jax.tree_util.tree_unflatten(tree_structure, outvals)
# 
#   return outvals, env


from typing import Optional
from jax_verify.src import bound_propagation
import jax
import jax.numpy as jnp
  """Wraps a Bound with additional concrete constraints."""

  def __init__(
      self,
      base_bound: bound_propagation.Bound,
      lower: Optional[Tensor],
      upper: Optional[Tensor],
  ):
    self._base_bound = base_bound
    self._lower = lower
    self._upper = upper

  def unwrap(self) -> bound_propagation.Bound:
    return self._base_bound.unwrap()

  @property
  def lower(self) -> Tensor:
    if self._lower is None:
      return self._base_bound.lower
    else:
      return jnp.maximum(self._base_bound.lower, self._lower)

  @property
  def upper(self) -> Tensor:
    if self._upper is None:
      return self._base_bound.upper
    else:
      return jnp.minimum(self._base_bound.upper, self._upper)



# Path: jax_verify/src/bound_propagation.py
# def bound_propagation(
#     prop_alg: PropagationAlgorithm[Repr],
#     function: Callable[..., Nest[Tensor]],
#     *bounds: Nest[GraphInput],
#     graph_simplifier=synthetic_primitives.default_simplifier,
# ) -> Tuple[
#     Nest[Union[Repr, Tensor]],
#     Dict[jax.core.Var, Union[Repr, Tensor, Bound]]]:
#   """Performs Bound Propagation on the model implemented by `function`.
# 
#   Args:
#     prop_alg: Algorithm specifying how to traverse the graph and how to
#       transform each node.
#     function: Pure function inputs -> outputs. If the function to propagate
#       through has a more complex signature, the use of `functools.partial` can
#       solve that problem.
#     *bounds: Nest of `IntervalBound` objects containing the lower and upper
#       bounds on all the inputs, or `Tensor`s containing known inputs directly.
#     graph_simplifier: Function transforming the JaxPR graph into a simpler
#       graph. Default value is a function identifying specific activation
#       functions, followed by grouping of linear sequences and quadratic forms.
#   Returns:
#     bounds: Bounds over all the outputs of the function, with the same structure
#       as the output of `function`
#     env: Mapping from the node of the computations to their representation.
#   """
#   # Replace all the jittable bounds by standard bound object.
#   bounds = unjit_inputs(*bounds)
# 
#   # Parse the computation graph.
#   placeholder_inputs = jax.tree_util.tree_map(
#       lambda b: b.lower if isinstance(b, Bound) else b,
#       bounds)
#   jaxpr_maker = jax.make_jaxpr(function)
#   parsed = jaxpr_maker(*placeholder_inputs)
#   output_shapes = jax.eval_shape(function, *placeholder_inputs)
# 
#   flat_is_bound, _ = jax.tree_util.tree_flatten(
#       jax.tree_util.tree_map(lambda b: isinstance(b, Bound), bounds))
#   inp_is_bound = {var: is_bound
#                   for var, is_bound in zip(parsed.jaxpr.invars, flat_is_bound)}
#   simplified_graph = synthetic_primitives.simplify_graph(
#       graph_simplifier, parsed.jaxpr, inp_is_bound)
#   graph = PropagationGraph(simplified_graph, parsed.literals)
# 
#   outvals, env = prop_alg.propagate(graph, bounds)
# 
#   # Make outvals into the same tree structure than the output of the function.
#   tree_structure = jax.tree_util.tree_structure(output_shapes)
#   outvals = jax.tree_util.tree_unflatten(tree_structure, outvals)
# 
#   return outvals, env


# Path: jax_verify/src/intersection.py
from typing import Optional
from jax_verify.src import bound_propagation
import jax
import jax.numpy as jnp

  """Wraps a Bound with additional concrete constraints."""

  def __init__(
      self,
      base_bound: bound_propagation.Bound,
      lower: Optional[Tensor],
      upper: Optional[Tensor],
  ):
    self._base_bound = base_bound
    self._lower = lower
    self._upper = upper

  def unwrap(self) -> bound_propagation.Bound:
    return self._base_bound.unwrap()

  @property
  def lower(self) -> Tensor:
    if self._lower is None:
      return self._base_bound.lower
    else:
      return jnp.maximum(self._base_bound.lower, self._lower)

  @property
  def upper(self) -> Tensor:
    if self._upper is None:
      return self._base_bound.upper
    else:
      return jnp.minimum(self._base_bound.upper, self._upper)



class IntersectionBoundTransform(bound_propagation.BoundTransform):