Trivial change for backward compatibility stub.

PiperOrigin-RevId: 445564091
poorse · Apr 30, 2022 · 290c90d · 290c90d
1 parent a682a37
commit 290c90d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
@@ -2135,3 +2135,5 @@ def trace_to_subjaxpr(main: core.MainTrace, instantiate: Union[bool, Sequence[bo
   out_pvals = [t.pval for t in out_tracers]
   del trace, in_tracers, out_tracers
   yield jaxpr, (out_pvals, consts, env)
+
+partial_eval_jaxpr: Callable