rocm_jax/jax/interpreters/partial_eval.py

# Copyright 2018 Google LLC
#
# 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
#
#     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.

import itertools as it
from collections import namedtuple
import contextlib
import functools
from typing import (Any, Callable, Dict, NamedTuple, Optional, Sequence, Tuple,
                    List, Union, cast, Type, no_type_check)
from weakref import ref

import numpy as np

from .. import core
from .. import dtypes
from .. import linear_util as lu
from ..abstract_arrays import ConcreteArray, raise_to_shaped
from ..ad_util import Zero
from ..util import (unzip2, safe_zip, safe_map, toposort, partial, split_list,
                    cache)
from ..core import (Trace, Tracer, Jaxpr, Literal, get_aval, AbstractValue,
                    unit, unitvar, abstract_unit, TypedJaxpr, new_jaxpr_eqn,
                    dropvar)
from .. import source_info_util
from ..config import config

map = safe_map
zip = safe_zip
def identity(x): return x

class PartialVal(tuple):
  """Partial value: either a known value or an unknown (abstract) value.

  Represented as a pair `(aval_opt, const)` of one of two kinds:
  * `(None, <Constant>)` indicates a known value, either a Python regular
    value, or a Tracer.
  * `(<AbstractValue>, *)` indicates an unknown value characterized by an
    abstract value.
  """
  def __new__(cls, xs: Tuple[Optional[AbstractValue], core.Value]):
    pv, const = xs
    if not core.skip_checks:
      # type checks
      assert isinstance(pv, (AbstractValue, type(None))), xs
      assert isinstance(const, core.Tracer) or type(const) is Zero or core.valid_jaxtype(const), xs
      # invariant checks
      if isinstance(pv, AbstractValue):
        assert get_aval(const) == core.abstract_unit, xs
    return tuple.__new__(cls, xs)

  @classmethod
  def known(cls, const: core.Value) -> 'PartialVal':
    return PartialVal((None, const))

  @classmethod
  def unknown(cls, aval: AbstractValue) -> 'PartialVal':
    return PartialVal((aval, core.unit))

  def is_known(self) -> bool:
    return self[0] is None

  def get_known(self) -> Optional[core.Value]:
    """Get the known value, if known, else None."""
    return self[1] if self[0] is None else None

  def get_aval(self) -> AbstractValue:
    """Get AbstractValue directly (if unknown) or from the constant (known)."""
    known = self.get_known()
    if known is not None:
      return get_aval(known)
    else:
      return self[0]

  def merge_with_known(self, val: core.Value) -> core.Value:
    """Either the stored known value, or the given 'val'."""
    known = self.get_known()
    return known if known is not None else val


class JaxprTrace(Trace):
  def pure(self, val) -> 'JaxprTracer':
    return self.new_const(val)

  def lift(self, val) -> 'JaxprTracer':
    return self.new_const(val)

  def sublift(self, val) -> 'JaxprTracer':
    return JaxprTracer(self, val.pval, FreeVar(val))

  def new_const(self, val) -> 'JaxprTracer':
    if isinstance(val, Tracer) and val._trace.level == self.level:
      raise Exception
    return JaxprTracer(self, PartialVal.known(val), unit)

  def new_instantiated_literal(self, val) -> 'JaxprTracer':
    return JaxprTracer(self, PartialVal.unknown(get_aval(val)), Literal(val))

  def new_instantiated_const(self, val) -> 'JaxprTracer':
    return JaxprTracer(self, PartialVal.unknown(get_aval(val)), ConstVar(val))

  def new_arg(self, pval: PartialVal) -> 'JaxprTracer':
    const = pval.get_known()
    if const is None:
      return JaxprTracer(self, pval, LambdaBinding())
    else:
      return self.new_const(const)

  def instantiate_const(self, tracer) -> Tracer:
    const = tracer.pval.get_known()
    if const is None:
      return tracer
    else:
      if type(const) in core.literalable_types and np.shape(const) == ():
        return self.new_instantiated_literal(const)
      else:
        return self.new_instantiated_const(const)

  def instantiate_const_abstracted(self, tracer) -> 'JaxprTracer':
    const = tracer.pval.get_known()
    if const is None:
      return tracer
    else:
      aval = raise_to_shaped(get_aval(const), np.isscalar(const))
      return JaxprTracer(self, PartialVal.unknown(aval), ConstVar(const))

  def process_primitive(self, primitive, tracers, params):
    if primitive in custom_partial_eval_rules:
      return custom_partial_eval_rules[primitive](self, *tracers, **params)
    else:
      return self.default_process_primitive(primitive, tracers, params)

  def default_process_primitive(self, primitive, tracers, params):
    """By default, if all the input tracers are known, then execute the primitive
    and all the ouputs are known. Otherwise, all the outputs are unknown."""
    consts = [t.pval.get_known() for t in tracers]
    if all(c is not None for c in consts):
      return primitive.bind(*consts, **params)
    tracers = map(self.instantiate_const, tracers)
    avals = [t.aval for t in tracers]
    out_aval = primitive.abstract_eval(*avals, **params)
    source = source_info_util.current()
    if primitive.multiple_results:
      out_tracers = [JaxprTracer(self, PartialVal.unknown(aval), None)
                     for aval in out_aval]
      eqn = new_eqn_recipe(tracers, out_tracers, primitive, params, source)
      for t in out_tracers: t.recipe = eqn
      return out_tracers
    else:
      out_tracer = JaxprTracer(self, PartialVal.unknown(out_aval), None)
      out_tracer.recipe = new_eqn_recipe(tracers, [out_tracer], primitive,
                                         params, source)
      return out_tracer

  # We use process_call to handle both call and map primitives.
  def process_call(self, primitive, f: lu.WrappedFun, tracers, params):
    if not config.omnistaging_enabled:
      if (self.main.trace_type is StagingJaxprTrace  # type: ignore
          and primitive in staged_out_calls):        # type: ignore
        tracers = map(self.instantiate_const_abstracted, tracers)

    if primitive in call_partial_eval_rules:
      return call_partial_eval_rules[primitive](self, primitive, f, tracers, params)

    in_pvals = [t.pval for t in tracers]
    if primitive.map_primitive:
      mapped_aval = partial(core.mapped_aval, params['axis_size'])
      in_pvals = [pval if pval.is_known() or not is_mapped
                  else PartialVal.unknown(mapped_aval(pval[0]))
                  for pval, is_mapped in zip(in_pvals, params['mapped_invars'])]
    jaxpr, out_pvals, consts, env_tracers = self.partial_eval(
        f, in_pvals, partial(primitive.bind, **params))
    if primitive.map_primitive:
      unmapped_aval = partial(core.unmapped_aval, params['axis_size'])
      out_pvals = [pval if pval.is_known()
                   else PartialVal.unknown(unmapped_aval(pval[0]))
                   for pval in out_pvals]

    # Avoid staging out trivial calls, but maps may involve broadcasting.
    if not jaxpr.eqns and not primitive.map_primitive:
      env = {core.unitvar: core.unit}
      map(env.setdefault, jaxpr.invars, (*env_tracers, *tracers))
      map(env.setdefault, jaxpr.constvars, consts)
      return [v.val if type(v) is Literal
              else pval.get_known() if pval.is_known()
              else env[v]
              for v, pval in zip(jaxpr.outvars, out_pvals)]

    # Skip known invars and outvars, and lift constants as regular invars
    in_knowns = tuple(t.pval.is_known() for t in it.chain(env_tracers, tracers))
    out_unknowns = tuple(not pval.is_known() for pval in out_pvals)
    jaxpr = _drop_invars(jaxpr, in_knowns)
    jaxpr = _dce_untyped_jaxpr(jaxpr, out_unknowns, drop_outputs=True)

    # Known tracers get propagated as if they were constants
    known_tracers_out = [self.new_const(pval.get_known()) for pval in out_pvals
                         if pval.is_known()]

    # Unknown tracers need to have the jaxpr set up as their recipe
    unknown_tracers_out = [JaxprTracer(self, pval, None) for pval in out_pvals
                           if not pval.is_known()]
    unknown_tracers_in = [t for t in tracers if not t.pval.is_known()]
    const_tracers = map(self.new_instantiated_const, consts)
    in_tracers = (*const_tracers, *env_tracers, *unknown_tracers_in)

    # Set up new params
    new_params = dict(params, call_jaxpr=convert_constvars_jaxpr(jaxpr))
    if primitive.map_primitive:
      mapped_invars = params['mapped_invars']
      new_mapped_invars = ((True,) * len(const_tracers) +
                           (False,) * len(env_tracers) +
                           tuple(v for v, t in zip(mapped_invars, tracers)
                                 if not t.pval.is_known()))
      new_params = dict(new_params, mapped_invars=new_mapped_invars)
    update_params = call_param_updaters.get(primitive)
    if update_params:
      new_params = update_params(new_params, [not t.pval.is_known() for t in tracers])

    eqn = new_eqn_recipe(in_tracers, unknown_tracers_out, primitive, new_params,
                         source_info_util.current())
    for t in unknown_tracers_out: t.recipe = eqn
    return _zip_knowns(known_tracers_out, unknown_tracers_out, out_unknowns)

  process_map = process_call

  # We use post_process_call to handle both call and map primitives.
  def post_process_call(self, primitive, out_tracers, params):
    jaxpr, consts, env = tracers_to_jaxpr([], out_tracers)
    out_pvs, out_pv_consts = unzip2(t.pval for t in out_tracers)
    out = out_pv_consts + consts
    del consts, out_pv_consts
    main = self.main

    if primitive.map_primitive:
      sz = params['axis_size']
      out_pvs = [None if pv is None else core.unmapped_aval(sz, pv)
                 for pv in out_pvs]

    def todo(x):
      n = len(jaxpr.outvars)
      out_pv_consts, consts = x[:n], x[n:]
      trace = JaxprTrace(main, core.cur_sublevel())
      const_tracers = map(trace.new_instantiated_const, consts)
      out_tracers = [JaxprTracer(trace, PartialVal((out_pv, out_pv_const)), None)
                     for out_pv, out_pv_const in zip(out_pvs, out_pv_consts)]
      in_tracers = (*const_tracers, *map(trace.full_raise, env))

      new_params = dict(params, call_jaxpr=convert_constvars_jaxpr(jaxpr))
      if primitive.map_primitive:
        new_mapped_invars = (True,) * len(const_tracers) + (False,) * len(env)
        new_params = dict(new_params, mapped_invars=new_mapped_invars)
      update_params = call_param_updaters.get(primitive)
      if update_params:
        new_params = update_params(new_params, [])

      eqn = new_eqn_recipe(in_tracers, out_tracers, primitive, new_params,
                           source_info_util.current())
      for t in out_tracers:
        t.recipe = eqn
      return out_tracers
    return out, todo

  post_process_map = post_process_call

  def partial_eval(self, f: lu.WrappedFun, pvals: Sequence[PartialVal],
                   app: Callable[[lu.WrappedFun, Tuple[core.Value, ...]], Tuple[core.Value]]):
    """Partially evaluate f on a sequence of PartialVals."""
    in_avals, in_consts = unzip2(pvals)
    f = trace_to_subjaxpr(f, self.main, False)
    f, aux = partial_eval_wrapper(f, tuple(in_avals))
    out_flat, (out_avals, jaxpr, env) = app(f, *in_consts), aux()
    out_consts, consts = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])
    out_pvs = map(PartialVal, zip(out_avals, out_consts))
    env_tracers = map(self.full_raise, env)
    return jaxpr, out_pvs, consts, env_tracers


@lu.transformation_with_aux
def partial_eval_wrapper(pvs: Sequence[Optional[AbstractValue]], *consts):
  py_args = map(PartialVal, zip(pvs, consts))
  jaxpr, (out_pvals, consts, env) = yield (py_args,), {}
  out_pvs, out_consts = unzip2(out_pvals)
  out = tuple(out_consts) + tuple(consts)
  yield out, (out_pvs, jaxpr, env)


custom_partial_eval_rules: Dict[core.Primitive, Callable] = {}
call_partial_eval_rules: Dict[core.Primitive, Callable] = {}
call_param_updaters: Dict[core.Primitive, Callable] = {}


def abstract_eval_fun(fun, *avals, **params):
  if config.omnistaging_enabled:
    _, avals_out, _ = trace_to_jaxpr_dynamic(lu.wrap_init(fun, params), avals)
  else:
    pvals_in = [PartialVal.unknown(a) for a in avals]
    _, pvals_out, _ = trace_to_jaxpr(lu.wrap_init(fun, params), pvals_in,
                                    instantiate=True, stage_out=True)  # type: ignore
    avals_out, _ = unzip2(pvals_out)
  for aval_out in avals_out:
    assert isinstance(aval_out, AbstractValue)  # instantiate=True
  return avals_out


JaxprTracerRecipe = Union['JaxprEqnRecipe', 'LambdaBinding', 'FreeVar',
                          'ConstVar', Literal, core.Unit]

class JaxprTracer(Tracer):
  __slots__ = ['pval', 'recipe']

  def __init__(self, trace: JaxprTrace, pval: PartialVal,
               recipe: Optional[JaxprTracerRecipe]):
    assert isinstance(pval, PartialVal)
    pv, const = pval
    if isinstance(const, Tracer) and const._trace.level >= trace.level:
      raise core.escaped_tracer_error(
          "Tracer from a higher level: {} in trace {}".format(const, trace))
    self._trace = trace
    self.pval = pval
    self.recipe = recipe

  def __repr__(self):
    return 'Traced<{}:{}>'.format(self.aval, self._trace)

  @property
  def aval(self) -> AbstractValue:
    return self.pval.get_aval()

  @property
  def parents(self) -> Sequence['JaxprTracer']:
    if isinstance(self.recipe, JaxprEqnRecipe):
      return self.recipe.invars
    else:
      return []

  def full_lower(self):
    known = self.pval.get_known()
    if known is not None:
      return core.full_lower(known)
    else:
      return self

  def is_known(self):
      return self.pval.is_known()

# TODO(necula): this should return a TypedJaxpr
def trace_to_jaxpr(fun: lu.WrappedFun, pvals: Sequence[PartialVal],
                   instantiate: Union[bool, Sequence[bool]] = False,
                   ) -> Tuple[Jaxpr, Tuple[PartialVal, ...], Tuple[core.Value, ...]]:
  """Traces a function into a Jaxpr, given PartialVals for inputs.

  Returns (`jaxpr`, `out_pvals`, `consts`). The `jaxpr` contains only the
  computation that depends on unknown inputs. The `out_pvals` are the PartialVal
  for the outputs. The intermediate values that depend only on known inputs and
  are needed to compute the output of `jaxpr` are in `consts` and are passed in
  as the constvars of the `jaxpr`. The handling of the known outputs depends on
  `instantiate`.

  For example, given `fun` defined as follows::

    def fun(ki, ui):  # ki will be a known input in this example
      ka = ki + 2
      kb = ka + 3
      return (kb, ui + ka)

  with `ki` the known PartialVal `1.`, and `ui` an unknown PartialVal. The only
  computation that depends on unknown inputs is `ui + ka` and will be the only
  computation in the body of the `jaxpr`. This computation depends on the known
  intermediate value `ka`, which will be computed statically. Currently, such
  constants are either embedded in the Jaxpr if they are scalars, or passed as a
  constvar to `jaxpr`, and then the value of the actual constant will be in
  `consts`:

  When `instantiate=False` we get::

    jaxpr =
      { lambda ka ; ki ui.
        let c = add ui ka
        in (*, c) }   # known outputs are `*`
    out_pvals = [PartialVal.known(6), PartialVal.unknown(ShapedArray)]
    consts = [3]  # the constant for `ka`

  When `instantiate=True` we get::

    jaxpr =
      { lambda ka kb ; ki ui.
        let c = add ui ka
        in (kb, c) }   # known output are explicit
    out_pvals = [PartialVal.unknown(ConcreteArray(6)), PartialVal.unknown(ShapedArray)]
    consts = [3, 6]  # values for `ka` and `kb` constvars
  """
  with core.new_main(JaxprTrace) as main:
    fun = trace_to_subjaxpr(fun, main, instantiate)
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    assert not env
    del main

  return jaxpr, out_pvals, consts


@lu.transformation
def trace_to_subjaxpr(main: core.MainTrace, instantiate: Union[bool, Sequence[bool]],
                      pvals: Sequence[PartialVal]):
  assert all([isinstance(pv, PartialVal) for pv in pvals]), pvals
  trace = JaxprTrace(main, core.cur_sublevel())
  in_tracers = map(trace.new_arg, pvals)
  ans = yield in_tracers, {}
  instantiate = [instantiate] * len(ans) if isinstance(instantiate, bool) else instantiate
  out_tracers = map(trace.full_raise, map(core.full_lower, ans))
  out_tracers = map(partial(instantiate_const_at, trace), instantiate, out_tracers)
  jaxpr, consts, env = tracers_to_jaxpr(in_tracers, out_tracers)
  out_pvals = [t.pval for t in out_tracers]
  del trace, in_tracers, out_tracers
  yield jaxpr, (out_pvals, consts, env)

def instantiate_const_at(trace: JaxprTrace, instantiate: bool, tracer):
  if instantiate:
    return trace.instantiate_const(trace.full_raise(tracer))
  else:
    return tracer


FreeVar = namedtuple('FreeVar', ['val'])
ConstVar = namedtuple('ConstVar', ['val'])
LambdaBinding = namedtuple('LambdaBinding', [])
class JaxprEqnRecipe(NamedTuple):
  eqn_id: object
  invars: Sequence[JaxprTracer]
  outvars: 'Sequence[ref[JaxprTracer]]'
  primitive: core.Primitive
  params: Dict[str, Any]
  source_info: Optional[source_info_util.Traceback]

def new_eqn_recipe(invars: Sequence[JaxprTracer],
                   outvars: Sequence[JaxprTracer],
                   primitive: core.Primitive,
                   params: Dict[str, Any],
                   source_info: Optional[source_info_util.Traceback]
                  ) -> JaxprEqnRecipe:
  """Constructs a new JaxEqnRecipe.

  Params:
    invars: the tracers for the primitive inputs.
    outvars: the tracers for the primitive outputs.
    primitive: the primitive.
    params: the primitive params
  """
  # TODO(necula): move these checks to core.check_jaxpr, and call in more places
  if primitive.call_primitive or primitive.map_primitive:
    assert "call_jaxpr" in params
  if primitive.map_primitive:
    assert ("mapped_invars" in params and
            len(params["mapped_invars"]) == len(params["call_jaxpr"].invars))
    assert ("donated_invars" in params and
            len(params["donated_invars"]) == len(params["call_jaxpr"].invars))
  return JaxprEqnRecipe(object(), tuple(invars), map(ref, outvars), primitive,
                        params, source_info)


def recipe_to_eqn(getvar: Callable[[JaxprTracer], core.Atom],
                  recipe: JaxprEqnRecipe) -> core.JaxprEqn:
  _, in_tracers, out_tracer_refs, primitive, params, source_info = recipe
  out_tracers = [t_ref() for t_ref in out_tracer_refs]
  invars  = [getvar(t) for t in in_tracers]
  outvars = [core.dropvar if t is None else cast(core.Var, getvar(t))
             for t in out_tracers]
  return new_jaxpr_eqn(invars, outvars, primitive, params, source_info)

def tracers_to_jaxpr(
  in_tracers: List[JaxprTracer],
  out_tracers: List[JaxprTracer]
  ) -> Tuple[Jaxpr, Tuple[Any, ...], Tuple[Any, ...]]:
  """Constructs Jaxpr given tracers for inputs and outputs.

  Params:
    in_tracers: the tracers that were created for the function inputs
    out_tracers: the tracers that were output by the function.

  Returns: a triple of a `Jaxpr`, a list of constant values corresponding to
    the `constvars` in the returned Jaxps, and a list of environment values.
    The vars for the environment values have been prepended to the Jaxpr's
    `invars`.
  """
  newvar = core.gensym()
  t_to_var: Dict[int, core.Atom] = {}
  def getvar(t: JaxprTracer) -> core.Atom:
    var = t_to_var.get(id(t))
    if var is None:
      aval = t.pval.get_aval() if not t.pval.is_known() else abstract_unit
      var = t_to_var[id(t)] = newvar(aval)
    return var
  sorted_tracers = toposort(out_tracers)
  invars = map(getvar, in_tracers)
  eqns: List[core.JaxprEqn] = []
  env: Dict[core.Var, Any] = {}
  consts: Dict[core.Var, Any] = {}
  const_to_var: Dict[int, core.Var] = {}
  def getconstvar(c):
    var = const_to_var.get(id(c))
    if var is None:
      var = const_to_var[id(c)] = newvar(get_aval(c))
    return var
  processed_eqn_ids = set()
  for t in sorted_tracers:
    recipe = t.recipe
    if isinstance(recipe, JaxprEqnRecipe):
      if recipe.eqn_id not in processed_eqn_ids:
        eqns.append(recipe_to_eqn(getvar, recipe))
        processed_eqn_ids.add(recipe.eqn_id)
    elif isinstance(recipe, LambdaBinding):
      if not any(t is in_tracer for in_tracer in in_tracers):
        raise core.escaped_tracer_error(
            "Tracer not among input tracers {}".format(t))
      assert in_tracers, "Lambda binding with no args"
    elif isinstance(recipe, FreeVar):
      env[cast(core.Var, getvar(t))] = recipe.val
    elif isinstance(recipe, ConstVar):
      v = t_to_var[id(t)] = getconstvar(recipe.val)
      consts[v] = recipe.val
    elif isinstance(recipe, Literal):
      t_to_var[id(t)] = recipe
    elif recipe is unit:
      t_to_var[id(t)] = unitvar
    else:
      raise TypeError(recipe)

  env_vars, env_vals = unzip2(env.items())
  const_vars, const_vals = unzip2(consts.items())
  # The env_vars are pre-pended to the invars
  jaxpr = Jaxpr(const_vars, [*env_vars, *invars], map(getvar, out_tracers), eqns)
  core.skip_checks or core.check_jaxpr(jaxpr)
  return jaxpr, const_vals, env_vals

@cache()
def convert_constvars_jaxpr(jaxpr: Jaxpr):
  """Moves the constvars to the start of invars."""
  core.skip_checks or core.check_jaxpr(jaxpr)
  lifted_jaxpr = Jaxpr(constvars=(),
                       invars=jaxpr.constvars + jaxpr.invars,
                       outvars=jaxpr.outvars, eqns=jaxpr.eqns)
  core.skip_checks or core.check_jaxpr(lifted_jaxpr)
  return lifted_jaxpr

def _split_aval(unknown: bool, aval: AbstractValue) -> Tuple[AbstractValue, AbstractValue]:
  return (abstract_unit, aval) if unknown else (aval, abstract_unit)

def partial_eval_jaxpr(jaxpr: TypedJaxpr, unknowns: Sequence[bool],
                       instantiate: Union[bool, Sequence[bool]],
                       ) -> Tuple[TypedJaxpr, TypedJaxpr, Sequence[bool]]:
  """Specializes a Jaxpr given an indication of which inputs are known.

  Returns: (jaxpr_known, jaxpr_unknown, out_unknowns).

  `out_unknowns` specifies which outputs are unknown (depend on some unknown inputs).
  `jaxpr_known` takes the same inputs as `jaxpr`, ignores the unknown inputs,
  and performs *all* the computation in `jaxpr` that depends only on the known inputs.
  Outputs correspond to those of `jaxpr`, with the unknown ones replaced with `*`,
  appended with the known residuals (the intermediate computations in `jaxpr`
  that depend only on known inputs and that are needed to compute the unknown outputs).

  `jaxpr_unknown` takes the same inputs as `jaxpr` along with the known residuals
  computed by `jaxpr_known` and returns the same outputs as `jaxpr` with the known
  outputs replaced by `*`.

  Roughly, `jaxpr(ki, ui)` is decomposed assuming `ki` and `ui` are the known and respectively
  unknown inputs into:

    jaxpr(ki, ui) = let kout, _, kresidual = jaxpr_known(kin, *)
                    let _, uout = jaxpr_unknown(ki, ui, kresidual)
                    in (kout, uout)

  For example, if `jaxpr` is lambda ki, ui: let ka = ki + 2
                                            in (ki + 3, ui + ka)"
  then
    `jaxpr_known` = lambda ki, ui: let ka = ki + 2
                                    in (ki + 3, *, ka)
    'jaxpr_unknown` = lambda ki, ui, ka: (*, ui + ka)

  Note that if instantiate is True for a given output, then jaxpr_known always returns a
  unit in its place. So when instantiate is True, the expectation is the one doesn't
  run `jaxpr_known` for any of its outputs, but only to generate residuals that will allow
  to obtain the full outputs once `jaxpr_unknown` is ran. Outputs known ahead of time will
  simply get passed as residual constants and returned immediately.
  """
  f = lu.wrap_init(core.jaxpr_as_fun(jaxpr))

  cell = []
  def fun(*vals):
    pvals = [PartialVal.unknown(aval) if uk else PartialVal.known(val)
            for aval, val, uk in zip(jaxpr.in_avals, vals, unknowns)]
    jaxpr_2, out_pvals_2, consts_2 = trace_to_jaxpr(f, pvals, instantiate=instantiate)
    out_pvs_2, out_consts_2 = unzip2(out_pvals_2)
    cell.append((out_pvs_2, jaxpr_2, len(consts_2)))
    return out_consts_2 + consts_2

  # For jaxpr_known we pass core.unit for the unknown inputs, and known PartialVal for the
  # known inputs.
  in_avals = [abstract_unit if uk else a for a, uk in zip(jaxpr.in_avals, unknowns)]
  jaxpr_1, out_avals, consts_1 = trace_to_jaxpr_dynamic(lu.wrap_init(fun), in_avals)
  (out_pvs_2, jaxpr_2, num_res), = cell
  assert len(jaxpr_2.constvars) == num_res

  #   jaxpr :: a -> b
  # jaxpr_1 :: a1 -> [b1, res]
  # jaxpr_2 :: res | a2 -> b2
  # jaxpr_2 :: [a2, res] -> b2
  jaxpr_2 = convert_constvars_jaxpr(jaxpr_2)
  jaxpr_2.invars = jaxpr_2.invars[num_res:] + jaxpr_2.invars[:num_res]
  for var, unknown in zip(jaxpr_2.invars[:len(unknowns)], unknowns):
    if not unknown:
      var.aval = abstract_unit

  uk_out = [pv is not None for pv in out_pvs_2]

  in_avals_1, in_avals_2 = unzip2(map(_split_aval, unknowns, jaxpr.in_avals))
  out_avals_1, out_avals_2 = unzip2(map(_split_aval, uk_out, jaxpr.out_avals))
  # out_avals_1 and in_avals_2 need the residuals added
  res_avals = out_avals[len(jaxpr.out_avals):]
  assert len(res_avals) == num_res
  out_avals_1 = [*out_avals_1, *res_avals]
  in_avals_2 = [*in_avals_2, *res_avals]

  typed_jaxpr_1 = TypedJaxpr(jaxpr_1, consts_1, in_avals_1, out_avals_1)
  typed_jaxpr_2 = TypedJaxpr(jaxpr_2, (), in_avals_2, out_avals_2)
  return typed_jaxpr_1, typed_jaxpr_2, uk_out


remat_call_p = core.CallPrimitive('remat_call')
remat_call = remat_call_p.bind
remat_call_p.def_impl(core.call_impl)

def _remat_partial_eval(trace, _, f, tracers, params):
  concrete = params['concrete']

  # Unlike JaxprTrace.process_call, we want to form a jaxpr for the entirety of
  # the function being called, not just for the unknown parts. To do that, we
  # instantiate all the input tracers as constants in the jaxpr being formed.
  # Those tracers might have concrete avals, and doing abstract interpretation
  # on concrete avals engenders a tradeoff: it allows data-dependent Python
  # control flow to work, but it can in some cases lead to redundant FLOPs (done
  # both in the `bind` call below and the `core.jaxpr_as_fun` call). We use the
  # `concrete` parameter to switch this behavior, and if `concrete` is False
  # then we raise the avals to the Shaped level.
  if concrete:
    instantiated_tracers = map(trace.instantiate_const, tracers)
  else:
    instantiated_tracers = map(trace.instantiate_const_abstracted, tracers)

  # Using the instantiated tracers, run call_bind like JaxprTrace.process_call.
  in_pvals = [t.pval for t in instantiated_tracers]
  if config.omnistaging_enabled:
    jaxpr, eval_out_pvals, consts, env_tracers = trace.partial_eval(
      f, in_pvals, partial(remat_call_p.bind, **params))
  else:
    with core.initial_style_staging():  # type: ignore
      jaxpr, eval_out_pvals, consts, env_tracers = trace.partial_eval(
        f, in_pvals, partial(remat_call_p.bind, **params))

  # Convert consts to inputs, since they may contain Tracer instances.
  jaxpr = convert_constvars_jaxpr(jaxpr)
  const_tracers = map(trace.new_instantiated_const, consts)

  # Since we traced with everything marked as unknown, but we need to know which
  # outputs are known/unknown, we use partial_eval_jaxpr to get out_unknowns.
  in_avals = ([raise_to_shaped(t.pval.get_aval()) for t in const_tracers] +
              [raise_to_shaped(t.pval.get_aval()) for t in env_tracers] +
              [raise_to_shaped(pval.get_aval()) for pval in in_pvals])
  out_avals = [raise_to_shaped(abstract_unit if var is unitvar
                               else get_aval(var.val) if type(var) is Literal
                               else pval.get_aval())
               for var, pval in zip(jaxpr.outvars, eval_out_pvals)]
  typed_jaxpr = core.TypedJaxpr(jaxpr, (), in_avals, out_avals)
  in_unknowns = ([False] * len(consts) +
                 [not t.is_known() for t in it.chain(env_tracers, tracers)])
  if config.omnistaging_enabled:
    jaxpr_known, jaxpr_unknown, out_unknowns = partial_eval_jaxpr(
        typed_jaxpr, in_unknowns, instantiate=False)  # type: ignore
  else:
    jaxpr_known, jaxpr_unknown, out_unknowns = partial_eval_jaxpr(
        typed_jaxpr, in_unknowns, instantiate=False, trace_type=trace.main.trace_type)  # type: ignore
  out_knowns = [not b for b in out_unknowns]
  out_known_pvals, out_unknown_pvals = _partition_knowns(eval_out_pvals, out_unknowns)

  # Next, we need values for the outputs that should be known. Since consts
  # weren't passed through Python for evaluation, we need to evaluate jaxpr_known,
  # minus the residual outputs that we don't need. When `concrete=True`, as an
  # optimization we can avoid redoing *some* redundant FLOPs, namely those that
  # produced concrete avals at the output, simply by using those as computed
  # values. For the use case of inverse-mode ad in op-by-op ("eager mode")
  # evaluation, all the primal outputs should be concrete (thus not recomputed).
  to_compute = [type(pval[0]) is not ConcreteArray
                for uk, pval in zip(out_unknowns, eval_out_pvals) if not uk]
  num_outputs = len(jaxpr_unknown.out_avals)
  num_res = len(jaxpr_known.out_avals) - num_outputs
  jaxpr_known_nores = _dce_jaxpr(jaxpr_known, out_knowns + [False] * num_res, drop_outputs=True)
  jaxpr_known_comp = _dce_jaxpr(jaxpr_known_nores, to_compute)
  _, in_consts = unzip2(t.pval for t in it.chain(env_tracers, tracers))
  reconstructed_consts = core.jaxpr_as_fun(jaxpr_known_comp)(*consts, *in_consts)
  out_known_pvals = map(_reconstruct_pval, out_known_pvals, reconstructed_consts)

  # Known outputs should keep propagating as constants
  assert all(pv.is_known() for pv in out_known_pvals)
  known_output_tracers = [trace.new_const(pval.get_known())
                          for pval in out_known_pvals]
  # Unknown outputs get wrapped in tracers with the appropriate recipe
  unknown_output_tracers = [JaxprTracer(trace, out_pval, None)
                            for out_pval in out_unknown_pvals]

  # dce jaxpr outputs
  new_jaxpr = _dce_jaxpr(typed_jaxpr, out_unknowns, drop_outputs=True).jaxpr
  new_params = dict(params, call_jaxpr=new_jaxpr)

  # set up eqn for unknown outputs
  in_tracers = (*const_tracers, *env_tracers, *instantiated_tracers)
  eqn = new_eqn_recipe(in_tracers, unknown_output_tracers, remat_call_p, new_params,
                       source_info_util.current())
  for t in unknown_output_tracers: t.recipe = eqn
  return _zip_knowns(known_output_tracers, unknown_output_tracers, out_unknowns)
call_partial_eval_rules[remat_call_p] = _remat_partial_eval

def _partition_knowns(pvals, unknowns: Sequence[bool]):
  return ([e for e, unknown in zip(pvals, unknowns) if not unknown],
          [e for e, unknown in zip(pvals, unknowns) if unknown])

def _zip_knowns(known_list, unknown_list, which_unknown: Sequence[bool]):
  known_iter, unknown_iter = iter(known_list), iter(unknown_list)
  return [next(unknown_iter) if uk else next(known_iter) for uk in which_unknown]


def _dce_jaxpr(typed_jaxpr: TypedJaxpr, outputs: Sequence[bool], drop_outputs=False) -> TypedJaxpr:
  if drop_outputs:
    new_out_avals = [aval for aval, output in zip(typed_jaxpr.out_avals, outputs) if output]
  else:
    new_out_avals = [aval if output else core.abstract_unit
                     for aval, output in zip(typed_jaxpr.out_avals, outputs)]
  new_jaxpr = _dce_untyped_jaxpr(typed_jaxpr.jaxpr, tuple(outputs), drop_outputs)
  return core.TypedJaxpr(new_jaxpr, typed_jaxpr.literals, typed_jaxpr.in_avals,
                         new_out_avals)

@cache()
def _dce_untyped_jaxpr(jaxpr: Jaxpr, outputs: Tuple[bool, ...], drop_outputs=False) -> Jaxpr:
  # This dead-code elimination is pretty rudimentary, and in particular doesn't
  # nontrivially DCE through scan, call, or other higher-order primitives.
  # TODO(mattjj): better DCE
  if drop_outputs:
    new_outvars = [var for var, output in zip(jaxpr.outvars, outputs) if output]
  else:
    new_outvars = [var if output else unitvar
                   for var, output in zip(jaxpr.outvars, outputs)]

  needed_vars = {v for v in new_outvars if type(v) is not Literal}
  new_eqns = []
  for eqn in jaxpr.eqns[::-1]:
    if set(eqn.outvars) & needed_vars:
      new_eqns.append(eqn)
      needed_vars.update(v for v in eqn.invars if type(v) is not Literal)
  new_eqns = new_eqns[::-1]
  return core.Jaxpr(jaxpr.constvars, jaxpr.invars,
                    new_outvars, new_eqns)

@cache()
def _drop_invars(jaxpr: Jaxpr, drop: Tuple[bool, ...]):
  return core.Jaxpr(jaxpr.constvars, [v for v, d in zip(jaxpr.invars, drop) if not d],
                    jaxpr.outvars, jaxpr.eqns)


def _reconstruct_pval(pval1: PartialVal, const2: core.Value):
  pv1, _ = pval1
  if pval1.is_known():
    return pval1
  else:
    if type(pv1) is ConcreteArray:
      return PartialVal.known(pv1.val)  # pytype: disable=attribute-error
    else:
      return PartialVal.known(const2)


def move_binders_to_front(typed_jaxpr: TypedJaxpr, to_move: Sequence[bool]) -> TypedJaxpr:
  """Reorder the `invars` to move to front the ones for which `to_move` is True."""
  assert not typed_jaxpr.jaxpr.constvars
  assert len(typed_jaxpr.in_avals) == len(to_move)
  new_invars = _move_to_front(typed_jaxpr.jaxpr.invars, to_move)
  new_jaxpr = core.Jaxpr((), new_invars, typed_jaxpr.jaxpr.outvars,
                         typed_jaxpr.jaxpr.eqns)
  new_in_avals = _move_to_front(typed_jaxpr.in_avals, to_move)
  new_typed_jaxpr = core.TypedJaxpr(new_jaxpr, typed_jaxpr.literals,
                                    new_in_avals, typed_jaxpr.out_avals)
  return new_typed_jaxpr

def _move_to_front(lst: Sequence, to_move: Sequence[bool]) -> Sequence:
  return ([elt for elt, move in zip(lst, to_move) if move] +
          [elt for elt, move in zip(lst, to_move) if not move])


class DynamicJaxprTracer(core.Tracer):
  __slots__ = ['aval', 'line_info']

  def __init__(self, trace, aval, line_info=None):
    self._trace = trace
    self.aval = aval
    self.line_info = line_info

  def full_lower(self):
    return self

  def _contents(self):
    return ()

  def _origin_msg(self):
    progenitor_eqns = self._trace.frame.find_progenitors(self)
    msgs = [f"  operation {core.pp_eqn(eqn, print_shapes=True)}\n"
            f"    from line {source_info_util.summarize(eqn.source_info)}"
            for eqn in progenitor_eqns]
    if msgs:
      origin = (f"While tracing the function {self._trace.main.source_info}, "
                "this value became a tracer due to JAX operations on these lines:"
                "\n\n" + "\n\n".join(msgs))
    else:
      origin = ("The error occured while tracing the function "
                f"{self._trace.main.source_info}.")
    return origin

  def _assert_live(self) -> None:
    if not self._trace.main.jaxpr_stack:  # type: ignore
      msg = f"tracer created on line {source_info_util.summarize(self.line_info)}"
      raise core.escaped_tracer_error(msg)

class JaxprStackFrame:
  __slots__ = ['newvar', 'tracer_to_var', 'constid_to_var', 'constvar_to_val',
               'tracers', 'eqns']

  def __init__(self):
    self.newvar = core.gensym()
    self.tracer_to_var = {}
    self.constid_to_var = {}
    self.constvar_to_val = {}
    self.tracers = []   # circ refs, frame->tracer->trace->main->frame,
    self.eqns = []      # cleared when we pop frame from main

  def to_jaxpr(self, in_tracers, out_tracers):
    invars = [self.tracer_to_var[id(t)] for t in in_tracers]
    outvars = [self.tracer_to_var[id(t)] for t in out_tracers]
    constvars, constvals = unzip2(self.constvar_to_val.items())
    jaxpr = Jaxpr(constvars, invars, outvars, self.eqns)
    jaxpr, constvals = _inline_literals(jaxpr, constvals)
    # core.skip_checks or core.check_jaxpr(jaxpr)
    out_avals = [t.aval for t in out_tracers]
    return jaxpr, out_avals, constvals

  def find_progenitors(self, tracer):
    var = self.tracer_to_var.get(id(tracer))
    if not var:
      return []
    active_vars = {var}
    for eqn in self.eqns[::-1]:
      produced = set(eqn.outvars) & active_vars
      if produced:
        active_vars.difference_update(produced)
        active_vars.update(eqn.invars)
    return [eqn for eqn in self.eqns if set(eqn.invars) & active_vars]

def _inline_literals(jaxpr, constvals):
  consts = dict(zip(jaxpr.constvars, constvals))
  newvar = core.gensym()
  class var(dict):
    def __missing__(self, v):
      new_v = self[v] = newvar(v.aval)
      return new_v
  var = var()

  def lit(var: core.Var) -> Optional[Any]:
    val = consts.get(var)
    if type(val) in core.literalable_types and not np.shape(val):
      return Literal(val)
    else:
      return None

  used = {v for eqn in jaxpr.eqns for v in eqn.invars} | set(jaxpr.outvars)
  new_constvars = [var[v] for v in jaxpr.constvars if not lit(v)]
  new_constvals = [c for v, c in zip(jaxpr.constvars, constvals) if not lit(v)]
  new_invars = [var[v] for v in jaxpr.invars]
  new_eqns = [new_jaxpr_eqn([lit(v) or var[v] for v in eqn.invars],
                            [var[v] if v in used else dropvar for v in eqn.outvars],
                            eqn.primitive, eqn.params, eqn.source_info)
              for eqn in jaxpr.eqns]
  new_outvars = [lit(v) or var[v] for v in jaxpr.outvars]
  new_jaxpr = Jaxpr(new_constvars, new_invars, new_outvars, new_eqns)
  return new_jaxpr, new_constvals

class DynamicJaxprTrace(core.Trace):
  __slots__ = []  # type: ignore

  @property
  def frame(self):
    return self.main.jaxpr_stack[-1]  # pytype: disable=attribute-error

  def new_arg(self, aval):
    tracer = DynamicJaxprTracer(self, aval, source_info_util.current())
    self.frame.tracers.append(tracer)
    self.frame.tracer_to_var[id(tracer)] = self.frame.newvar(aval)
    return tracer

  def new_const(self, val):
    aval = raise_to_shaped(get_aval(val), weak_type=dtypes.is_python_scalar(val))
    tracer = DynamicJaxprTracer(self, aval, source_info_util.current())
    self.frame.tracers.append(tracer)
    var = self.frame.tracer_to_var[id(tracer)] = self.getconstvar(val)
    self.frame.constvar_to_val[var] = val
    return tracer

  pure = lift = sublift = new_const

  def getvar(self, tracer):
    var = self.frame.tracer_to_var.get(id(tracer))
    if var is None:
      self.frame.tracers.append(tracer)
      var = self.frame.tracer_to_var[id(tracer)] = self.frame.newvar(tracer.aval)
    return var

  def getconstvar(self, c):
    var = self.frame.constid_to_var.get(id(c))
    if var is None:
      var = self.frame.constid_to_var[id(c)] = self.frame.newvar(get_aval(c))
    return var

  def instantiate_const(self, val):
    if (isinstance(val, Tracer) and val._trace.main is self.main
        and val._trace.sublevel == self.sublevel):
      return val
    else:
      return self.new_const(val)

  def process_primitive(self, primitive, tracers, params):
    avals = [t.aval for t in tracers]
    out_avals = primitive.abstract_eval(*avals, **params)
    out_avals = [out_avals] if not primitive.multiple_results else out_avals
    source_info = source_info_util.current()
    out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
    invars = map(self.getvar, tracers)
    outvars = map(self.getvar, out_tracers)
    eqn = new_jaxpr_eqn(invars, outvars, primitive, params, source_info)
    self.frame.eqns.append(eqn)
    return out_tracers if primitive.multiple_results else out_tracers.pop()

  def process_call(self, call_primitive, f, tracers, params):
    in_avals = [t.aval for t in tracers]
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
    if not jaxpr.eqns:
      return core.eval_jaxpr(jaxpr, consts, *tracers)
    source_info = source_info_util.current()
    out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
    invars = map(self.getvar, tracers)
    outvars = map(self.getvar, out_tracers)
    constvars = map(self.getvar, map(self.instantiate_const, consts))
    new_params = dict(params, call_jaxpr=convert_constvars_jaxpr(jaxpr))
    update_params = call_param_updaters.get(call_primitive)
    if update_params:
      new_params = update_params(new_params, [True] * len(tracers))
    eqn = new_jaxpr_eqn([*constvars, *invars], outvars, call_primitive,
                        new_params, source_info)
    self.frame.eqns.append(eqn)
    return out_tracers

  def post_process_call(self, call_primitive, out_tracers, params):
    assert False  # unreachable

  def process_map(self, map_primitive, f, tracers, params):
    in_avals = [t.aval for t in tracers]
    axis_name, axis_size = params['axis_name'], params['axis_size']
    reduced_in_avals = [core.mapped_aval(axis_size, a) if m else a
                        for m, a in zip(params['mapped_invars'], in_avals)]
    with core.extend_axis_env(axis_name, axis_size, None):  # type: ignore
      jaxpr, reduced_out_avals, consts = trace_to_subjaxpr_dynamic(
          f, self.main, reduced_in_avals)
    out_avals = [core.unmapped_aval(params['axis_size'], a) for a in reduced_out_avals]
    source_info = source_info_util.current()
    out_tracers = [DynamicJaxprTracer(self, a, source_info) for a in out_avals]
    invars = map(self.getvar, tracers)
    outvars = map(self.getvar, out_tracers)
    constvars = map(self.getvar, map(self.instantiate_const, consts))
    new_mapped_invars = (False,) * len(consts) + params['mapped_invars']
    new_params = dict(params, mapped_invars=new_mapped_invars,
                      call_jaxpr=convert_constvars_jaxpr(jaxpr))
    update_params = call_param_updaters.get(map_primitive)
    if update_params:
      new_params = update_params(new_params, [True] * len(tracers))
    eqn = new_jaxpr_eqn([*constvars, *invars], outvars, map_primitive,
                        new_params, source_info)
    self.frame.eqns.append(eqn)
    return out_tracers

  def post_process_map(self, map_primitive, out_tracers, params):
    assert False  # unreachable


def trace_to_jaxpr_dynamic(fun: lu.WrappedFun, in_avals: Sequence[AbstractValue]):
  assert config.omnistaging_enabled
  with core.new_main(DynamicJaxprTrace, dynamic=True) as main:  # type: ignore
    main.source_info = fun_sourceinfo(fun.f)  # type: ignore
    main.jaxpr_stack = ()  # type: ignore
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
    del main
  return jaxpr, out_avals, consts

def trace_to_subjaxpr_dynamic(fun: lu.WrappedFun, main: core.MainTrace,
                              in_avals: Sequence[AbstractValue]):
  frame = JaxprStackFrame()
  with extend_jaxpr_stack(main, frame):
    trace = DynamicJaxprTrace(main, core.cur_sublevel())
    in_tracers = map(trace.new_arg, in_avals)
    ans = fun.call_wrapped(*in_tracers)
    out_tracers = map(trace.full_raise, ans)
  jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)
  return jaxpr, out_avals, consts

@contextlib.contextmanager
def extend_jaxpr_stack(main, frame):
  main.jaxpr_stack = main.jaxpr_stack + (frame,)
  try:
    yield
  finally:
    assert frame is main.jaxpr_stack[-1]
    main.jaxpr_stack = main.jaxpr_stack[:-1]

def trace_to_jaxpr_final(fun: lu.WrappedFun, in_avals: Sequence[AbstractValue]):
  assert config.omnistaging_enabled
  with core.new_base_main(DynamicJaxprTrace) as main:  # type: ignore
    main.source_info = fun_sourceinfo(fun.f)  # type: ignore
    main.jaxpr_stack = ()  # type: ignore
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
    del main
  return jaxpr, out_avals, consts

def partial_eval_to_jaxpr_dynamic(fun: lu.WrappedFun, in_pvals: Sequence[PartialVal]):
  # This function provides a partial evaluation behavior used by Flax. We can't
  # use trace_to_jaxpr directly because of an interaction with the curent
  # custom_derivatives.py, which we work around by adding the EvalTrace.
  # TODO(mattjj): alias to trace_to_jaxpr after revising custom_derivatives.py
  assert config.omnistaging_enabled
  with core.new_main(core.EvalTrace, dynamic=True) as _:  # type: ignore
    return trace_to_jaxpr(fun, in_pvals)

def fun_sourceinfo(fun):
  if isinstance(fun, functools.partial):
    fun = fun.func
  try:
    filename = fun.__code__.co_filename
    lineno = fun.__code__.co_firstlineno
    return f"{fun.__name__} at {filename}:{lineno}"
  except AttributeError:
    return "<unknown>"


@config.register_omnistaging_disabler
@no_type_check
def omnistaging_disabler() -> None:
  global trace_to_jaxpr, partial_eval_jaxpr, staged_out_calls, StagingJaxprTrace

  def trace_to_jaxpr(fun: lu.WrappedFun, pvals: Sequence[PartialVal],
                     instantiate: Union[bool, Sequence[bool]] = False,
                     stage_out=False, bottom=False,
                     trace_type: Optional[Type[Trace]] = None,
                     ) -> Tuple[Jaxpr, Tuple[PartialVal, ...], Tuple[core.Value, ...]]:
    """Traces a function into a Jaxpr, given PartialVals for inputs.

    Returns (`jaxpr`, `out_pvals`, `consts`). The `jaxpr` contains only the
    computation that depends on unknown inputs. The `out_pvals` are the PartialVal
    for the outputs. The intermediate values that depend only on known inputs and
    are needed to compute the output of `jaxpr` are in `consts` and are passed in
    as the constvars of the `jaxpr`. The handling of the known outputs depends on
    `instantiate`.

    For example, given `fun` defined as follows::

      def fun(ki, ui):  # ki will be a known input in this example
        ka = ki + 2
        kb = ka + 3
        return (kb, ui + ka)

    with `ki` the known PartialVal `1.`, and `ui` an unknown PartialVal. The only
    computation that depends on unknown inputs is `ui + ka` and will be the only
    computation in the body of the `jaxpr`. This computation depends on the known
    intermediate value `ka`, which will be computed statically. Currently, such
    constants are either embedded in the Jaxpr if they are scalars, or passed as a
    constvar to `jaxpr`, and then the value of the actual constant will be in
    `consts`:

    When `instantiate=False` we get::

      jaxpr =
        { lambda ka ; ki ui.
          let c = add ui ka
          in (*, c) }   # known outputs are `*`
      out_pvals = [PartialVal.known(6), PartialVal.unknown(ShapedArray)]
      consts = [3]  # the constant for `ka`

    When `instantiate=True` we get::

      jaxpr =
        { lambda ka kb ; ki ui.
          let c = add ui ka
          in (kb, c) }   # known output are explicit
      out_pvals = [PartialVal.unknown(ConcreteArray(6)), PartialVal.unknown(ShapedArray)]
      consts = [3, 6]  # values for `ka` and `kb` constvars
    """
    trace_type = trace_type or (StagingJaxprTrace if stage_out else JaxprTrace)
    with core.new_main(trace_type, bottom=bottom) as main:
      fun = trace_to_subjaxpr(fun, main, instantiate)
      jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
      assert not env
      del main

    return jaxpr, out_pvals, consts

  def partial_eval_jaxpr(jaxpr: TypedJaxpr, unknowns: Sequence[bool],
                        instantiate: Union[bool, Sequence[bool]],
                        trace_type: Optional[Type[core.Trace]]
                        ) -> Tuple[TypedJaxpr, TypedJaxpr, Sequence[bool]]:
    f = lu.wrap_init(core.jaxpr_as_fun(jaxpr))

    cell = []
    def fun(*vals):
      pvals = [PartialVal.unknown(aval) if uk else PartialVal.known(val)
              for aval, val, uk in zip(jaxpr.in_avals, vals, unknowns)]
      jaxpr_2, out_pvals_2, consts_2 = trace_to_jaxpr(f, pvals, instantiate=instantiate,
                                                      trace_type=trace_type)
      out_pvs_2, out_consts_2 = unzip2(out_pvals_2)
      cell.append((out_pvs_2, jaxpr_2, len(consts_2)))
      return out_consts_2 + consts_2

    # The abstract_unit here doesn't really matter, because trace_to_jaxpr completely ignores
    # the avals, and it will never actually reach any primitives, because the `fun` above will
    # execute the jaxpr with the right avals (it reconstructs `pvals` inside).
    pvals = [PartialVal.unknown(abstract_unit) if uk else PartialVal.unknown(aval)
            for aval, uk in zip(jaxpr.in_avals, unknowns)]
    jaxpr_1, out_pvals, consts_1 = trace_to_jaxpr(lu.wrap_init(fun), pvals, instantiate=True)
    (out_pvs_2, jaxpr_2, num_res), = cell
    assert len(jaxpr_2.constvars) == num_res

    #   jaxpr :: a -> b
    # jaxpr_1 :: a1 -> [b1, res]
    # jaxpr_2 :: res | a2 -> b2
    # jaxpr_2 :: [a2, res] -> b2
    jaxpr_2 = convert_constvars_jaxpr(jaxpr_2)
    jaxpr_2.invars = jaxpr_2.invars[num_res:] + jaxpr_2.invars[:num_res]
    for var, unknown in zip(jaxpr_2.invars[:len(unknowns)], unknowns):
      if not unknown:
        var.aval = abstract_unit

    uk_out = [pv is not None for pv in out_pvs_2]

    in_avals_1, in_avals_2 = unzip2(map(_split_aval, unknowns, jaxpr.in_avals))
    out_avals_1, out_avals_2 = unzip2(map(_split_aval, uk_out, jaxpr.out_avals))
    # out_avals_1 and in_avals_2 need the residuals added
    out_pvs, _ = unzip2(out_pvals)
    res_avals = out_pvs[len(jaxpr.out_avals):]
    assert len(res_avals) == num_res
    out_avals_1 = out_avals_1 + res_avals
    in_avals_2 = in_avals_2 + res_avals

    typed_jaxpr_1 = TypedJaxpr(jaxpr_1, consts_1, in_avals_1, out_avals_1)
    typed_jaxpr_2 = TypedJaxpr(jaxpr_2, (), in_avals_2, out_avals_2)
    return typed_jaxpr_1, typed_jaxpr_2, uk_out

  def process_custom_jvp_call(self, prim, fun, jvp, tracers):
    # See comment at top of `JaxprTrace`. This method should be reachable
    # only when we stage out, and in that case we drop the custom differentiation
    # rules, because we do not need them.
    if not config.omnistaging_enabled:
      assert self.main.trace_type is StagingJaxprTrace
    return fun.call_wrapped(*tracers)
  JaxprTrace.process_custom_jvp_call = process_custom_jvp_call

  def process_custom_vjp_call(self, prim, fun, fwd, bwd, tracers, out_trees):
    # See comment in the above process_custom_jvp_call method.
    if not config.omnistaging_enabled:
      assert self.main.trace_type is StagingJaxprTrace
    return fun.call_wrapped(*tracers)
  JaxprTrace.process_custom_vjp_call = process_custom_vjp_call

  staged_out_calls = set()

  class StagingJaxprTrace(JaxprTrace): pass