# 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


import functools
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import jax_verify
import tree
from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.src import bound_propagation

    tight_input_bounds = jax_verify.IntervalBound(z, z)
    fun_to_prop = functools.partial(model.apply, params)
    tight_output_bounds = jax_verify.interval_bound_propagation(
        fun_to_prop, tight_input_bounds)

    model_eval = model.apply(params, z)

    # Because the input lower bound is equal to the input upper bound, the value
    # of the output bounds should be the same and correspond to the value of the
    # forward pass.
    self.assertAlmostEqual(tight_output_bounds.lower.tolist(),
                           tight_output_bounds.upper.tolist())

    self.assertAlmostEqual(tight_output_bounds.lower.tolist(),
                           model_eval.tolist())

  @parameterized.named_parameters(
      ('Sequential', sequential_model),
      ('Residual', residual_model),
      ('ResidualAll', residual_model_intermediate),
      ('1elt_list', single_element_list_model),
      ('dict_output', dict_output_model)
  )
  def test_matching_output_structure(self, model):
    def _check_matching_structures(output_tree, bound_tree):
      """Replace all bounds/arrays with True, then compare pytrees."""
      output_struct = tree.traverse(
          lambda x: True if isinstance(x, jnp.ndarray) else None, output_tree)
      bound_struct = tree.traverse(

# 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/tests/bound_propagation_test.py
import functools
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import jax_verify
import tree
from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.src import bound_propagation


    tight_input_bounds = jax_verify.IntervalBound(z, z)
    fun_to_prop = functools.partial(model.apply, params)
    tight_output_bounds = jax_verify.interval_bound_propagation(
        fun_to_prop, tight_input_bounds)

    model_eval = model.apply(params, z)

    # Because the input lower bound is equal to the input upper bound, the value
    # of the output bounds should be the same and correspond to the value of the
    # forward pass.
    self.assertAlmostEqual(tight_output_bounds.lower.tolist(),
                           tight_output_bounds.upper.tolist())

    self.assertAlmostEqual(tight_output_bounds.lower.tolist(),
                           model_eval.tolist())

  @parameterized.named_parameters(
      ('Sequential', sequential_model),
      ('Residual', residual_model),
      ('ResidualAll', residual_model_intermediate),
      ('1elt_list', single_element_list_model),
      ('dict_output', dict_output_model)
  )
  def test_matching_output_structure(self, model):
    def _check_matching_structures(output_tree, bound_tree):
      """Replace all bounds/arrays with True, then compare pytrees."""
      output_struct = tree.traverse(
          lambda x: True if isinstance(x, jnp.ndarray) else None, output_tree)
      bound_struct = tree.traverse(

          lambda x: True if isinstance(x, bound_propagation.Bound) else None,