Merge pull request #20768 from chaserileyroberts:chase/eval_jaxpr_in_extend

PiperOrigin-RevId: 625830536

Author: jax authors

jax/extend/core/__init__.py

@@ -19,4 +19,15 @@
   array_types as array_types
 )
 
+from jax._src.core import (
+  ClosedJaxpr as ClosedJaxpr,
+  Jaxpr as Jaxpr,
+  JaxprEqn as JaxprEqn,
+  jaxpr_as_fun as jaxpr_as_fun,
+  Literal as Literal,
+  Primitive as Primitive,
+  Token as Token,
+  Var as Var,
+)
+
 from . import primitives as primitives