# Path: jax_verify/src/synthetic_primitives.py
# T = TypeVar('T')
# LINEAR_OP = [
#     lax.reshape_p,
#     lax.squeeze_p,
#     lax.transpose_p,
#     lax.broadcast_in_dim_p,
#     lax.gather_p,
#     lax.reduce_sum_p,
#     lax.add_p,
#     lax.scatter_add_p,
#     lax.sub_p,
#     linear_p,
# ]
# BILINEAR_OP = [
#     lax.dot_general_p,
#     lax.conv_general_dilated_p,
#     lax.mul_p,
#     posbilinear_p,
# ]
# def simplify_graph(
#     graph_simplifier: GraphSimplifier,
#     graph: jax.core.Jaxpr,
#     var_is_bound: VarIsBoundDict,
# ) -> jax.core.Jaxpr:
# def _simplify_graph(
#     var_is_bound: VarIsBoundDict,
#     graph: jax.core.Jaxpr,
#     graph_simplifier: GraphSimplifier,
# ) -> jax.core.Jaxpr:
#   def __init__(
#       self,
#       fn: Callable[..., Tensor],
#       upstream_simplifier: Optional[SimpleSimplifier],
#       primitive: 'FakePrimitive',
#       *arg_shapes: Sequence[int],
#       **params):
#   def graph(self) -> jax.core.Jaxpr:
#   def capture_literals(self) -> Dict[int, str]:
#   def primitive(self) -> 'FakePrimitive':
#   def params(self) -> Dict[str, Any]:
#   def simplify(self, graph_simplifier: SimpleSimplifier):
# def _mark_outputs_whether_bounds(eqn, var_is_bound):
# def jax_primitive_subgraph(eqn: jax.core.JaxprEqn) -> Optional[jax.core.Jaxpr]:
# def _propagate_var_is_bound(graph: jax.core.Jaxpr,
#                             var_is_bound: VarIsBoundDict):
# def detect(
#     synthetic_primitive_specs: Sequence[SyntheticPrimitiveSpec],
#     graph: jax.core.Jaxpr,
# ) -> jax.core.Jaxpr:
# def _next_equation(
#     synthetic_primitive_specs: Sequence[SyntheticPrimitiveSpec],
#     graph: jax.core.Jaxpr,
#     eqn_idx: int,
# ) -> Tuple[jax.core.JaxprEqn, int]:
# def _equal_literal_values(lhs, rhs) -> bool:
# def _matches(
#     spec: jax.core.Jaxpr,
#     capture_literals: Dict[int, str],
#     graph: jax.core.Jaxpr,
#     eqn_idx: int,
# ) -> Tuple[
# def _differing_literals(
#     graph: jax.core.Jaxpr,
#     alt_graph: jax.core.Jaxpr,
# ) -> Sequence[jax.core.Literal]:
# def _is_linear_eqn(eqn: jax.core.JaxprEqn, var_is_bound: VarIsBoundDict):
# def _is_posbilinear_eqn(eqn, var_is_bound):
# def _find_eqn(eqn_list: List[jax.core.JaxprEqn], var: jax.core.Var) -> int:
# def group_posbilinear(graph: jax.core.Jaxpr,
#                       var_is_bound: VarIsBoundDict,
#                       ) -> jax.core.Jaxpr:
# def group_linear_sequence(graph: jax.core.Jaxpr,
#                           var_is_bound: VarIsBoundDict,
#                           ) -> jax.core.Jaxpr:
# def hoist_constant_computations(graph: jax.core.Jaxpr,
#                                 var_is_bound: VarIsBoundDict
#                                 ) -> jax.core.Jaxpr:
# def expand_softmax_simplifier(graph: jax.core.Jaxpr,
#                               var_is_bound: VarIsBoundDict
#                               ) -> jax.core.Jaxpr:
#   def __init__(self, name, impl):
#   def bind(self, *args, **kwargs):
#   def name(self):
#   def __str__(self):
#   def __repr__(self):
# def simplifier_composition(*graph_simplifiers: GraphSimplifier
#                            ) -> GraphSimplifier:
# def _subgraph_bind(*args, jax_verify_subgraph, jax_verify_keepjvargs):
#   def __init__(self, name):
#   def bind(self, *args, **kwargs):
# def activation_specs() -> Sequence[SyntheticPrimitiveSpec]:
# class SyntheticPrimitiveSpec:
# class FakePrimitive:
# class SubgraphPrimitive(FakePrimitive):


import abc
import collections
import dataclasses
import functools
import jax
import jax.numpy as jnp
from typing import Any, Callable, Dict, Generic, List, Optional, Sequence, Tuple, TypeVar, Union
from jax_verify.src import synthetic_primitives
    for bound, invar in zip(flat_bounds, self._graph.invars):
      if is_bound(bound):
        input_val = transform.input_transform(
            TransformContext(index.as_tuple(), None),
            bound.lower, bound.upper)
        self._index_to_node[index.as_tuple()] = invar
        index.incr()
      else:
        # Input is a fixed tensor.
        input_val = bound
      invals.append(input_val)
    env.update(zip(self._graph.invars, invals))
    env.update(zip(self._graph.constvars, self._literals))

    for eqn in self._graph.eqns:
      self._forward_prop_eqn(transform, env, index, eqn)
    self._last_index = index.as_tuple()

    outvals = jax.tree_map(functools.partial(read_env, env),
                           self._graph.outvars)
    return outvals, env

  def _forward_prop_eqn(self, transform: GraphTransform[Repr],
                        env: Dict[jax.core.Var, Union[Repr, Tensor, _Ref]],
                        index: IndexCounter,
                        eqn: jax.core.JaxprEqn):
    """Recursive step of `forward_propagation`."""
    def subgraph_handler(sub_transform, sub_graph, *invals):
      assert len(invals) == len(sub_graph.invars) == len(eqn.invars)
      use_ref = not isinstance(eqn.primitive,

# Path: jax_verify/src/synthetic_primitives.py
# T = TypeVar('T')
# LINEAR_OP = [
#     lax.reshape_p,
#     lax.squeeze_p,
#     lax.transpose_p,
#     lax.broadcast_in_dim_p,
#     lax.gather_p,
#     lax.reduce_sum_p,
#     lax.add_p,
#     lax.scatter_add_p,
#     lax.sub_p,
#     linear_p,
# ]
# BILINEAR_OP = [
#     lax.dot_general_p,
#     lax.conv_general_dilated_p,
#     lax.mul_p,
#     posbilinear_p,
# ]
# def simplify_graph(
#     graph_simplifier: GraphSimplifier,
#     graph: jax.core.Jaxpr,
#     var_is_bound: VarIsBoundDict,
# ) -> jax.core.Jaxpr:
# def _simplify_graph(
#     var_is_bound: VarIsBoundDict,
#     graph: jax.core.Jaxpr,
#     graph_simplifier: GraphSimplifier,
# ) -> jax.core.Jaxpr:
#   def __init__(
#       self,
#       fn: Callable[..., Tensor],
#       upstream_simplifier: Optional[SimpleSimplifier],
#       primitive: 'FakePrimitive',
#       *arg_shapes: Sequence[int],
#       **params):
#   def graph(self) -> jax.core.Jaxpr:
#   def capture_literals(self) -> Dict[int, str]:
#   def primitive(self) -> 'FakePrimitive':
#   def params(self) -> Dict[str, Any]:
#   def simplify(self, graph_simplifier: SimpleSimplifier):
# def _mark_outputs_whether_bounds(eqn, var_is_bound):
# def jax_primitive_subgraph(eqn: jax.core.JaxprEqn) -> Optional[jax.core.Jaxpr]:
# def _propagate_var_is_bound(graph: jax.core.Jaxpr,
#                             var_is_bound: VarIsBoundDict):
# def detect(
#     synthetic_primitive_specs: Sequence[SyntheticPrimitiveSpec],
#     graph: jax.core.Jaxpr,
# ) -> jax.core.Jaxpr:
# def _next_equation(
#     synthetic_primitive_specs: Sequence[SyntheticPrimitiveSpec],
#     graph: jax.core.Jaxpr,
#     eqn_idx: int,
# ) -> Tuple[jax.core.JaxprEqn, int]:
# def _equal_literal_values(lhs, rhs) -> bool:
# def _matches(
#     spec: jax.core.Jaxpr,
#     capture_literals: Dict[int, str],
#     graph: jax.core.Jaxpr,
#     eqn_idx: int,
# ) -> Tuple[
# def _differing_literals(
#     graph: jax.core.Jaxpr,
#     alt_graph: jax.core.Jaxpr,
# ) -> Sequence[jax.core.Literal]:
# def _is_linear_eqn(eqn: jax.core.JaxprEqn, var_is_bound: VarIsBoundDict):
# def _is_posbilinear_eqn(eqn, var_is_bound):
# def _find_eqn(eqn_list: List[jax.core.JaxprEqn], var: jax.core.Var) -> int:
# def group_posbilinear(graph: jax.core.Jaxpr,
#                       var_is_bound: VarIsBoundDict,
#                       ) -> jax.core.Jaxpr:
# def group_linear_sequence(graph: jax.core.Jaxpr,
#                           var_is_bound: VarIsBoundDict,
#                           ) -> jax.core.Jaxpr:
# def hoist_constant_computations(graph: jax.core.Jaxpr,
#                                 var_is_bound: VarIsBoundDict
#                                 ) -> jax.core.Jaxpr:
# def expand_softmax_simplifier(graph: jax.core.Jaxpr,
#                               var_is_bound: VarIsBoundDict
#                               ) -> jax.core.Jaxpr:
#   def __init__(self, name, impl):
#   def bind(self, *args, **kwargs):
#   def name(self):
#   def __str__(self):
#   def __repr__(self):
# def simplifier_composition(*graph_simplifiers: GraphSimplifier
#                            ) -> GraphSimplifier:
# def _subgraph_bind(*args, jax_verify_subgraph, jax_verify_keepjvargs):
#   def __init__(self, name):
#   def bind(self, *args, **kwargs):
# def activation_specs() -> Sequence[SyntheticPrimitiveSpec]:
# class SyntheticPrimitiveSpec:
# class FakePrimitive:
# class SubgraphPrimitive(FakePrimitive):


# Path: jax_verify/src/graph_traversal.py
import abc
import collections
import dataclasses
import functools
import jax
import jax.numpy as jnp
from typing import Any, Callable, Dict, Generic, List, Optional, Sequence, Tuple, TypeVar, Union
from jax_verify.src import synthetic_primitives

    for bound, invar in zip(flat_bounds, self._graph.invars):
      if is_bound(bound):
        input_val = transform.input_transform(
            TransformContext(index.as_tuple(), None),
            bound.lower, bound.upper)
        self._index_to_node[index.as_tuple()] = invar
        index.incr()
      else:
        # Input is a fixed tensor.
        input_val = bound
      invals.append(input_val)
    env.update(zip(self._graph.invars, invals))
    env.update(zip(self._graph.constvars, self._literals))

    for eqn in self._graph.eqns:
      self._forward_prop_eqn(transform, env, index, eqn)
    self._last_index = index.as_tuple()

    outvals = jax.tree_map(functools.partial(read_env, env),
                           self._graph.outvars)
    return outvals, env

  def _forward_prop_eqn(self, transform: GraphTransform[Repr],
                        env: Dict[jax.core.Var, Union[Repr, Tensor, _Ref]],
                        index: IndexCounter,
                        eqn: jax.core.JaxprEqn):
    """Recursive step of `forward_propagation`."""
    def subgraph_handler(sub_transform, sub_graph, *invals):
      assert len(invals) == len(sub_graph.invars) == len(eqn.invars)
      use_ref = not isinstance(eqn.primitive,

                               synthetic_primitives.FakePrimitive)