Re-export AbstractRef from hijax

Google-ML-Automation · Google-ML-Automation · commit e45b7b654c57 · 2025-11-20T12:37:07.000-08:00
PiperOrigin-RevId: 834868498
diff --git a/jax/experimental/BUILD b/jax/experimental/BUILD
@@ -513,6 +513,7 @@ pytype_strict_library(
         "//jax/_src:core",
         "//jax/_src:effects",
         "//jax/_src:hijax",
+        "//jax/_src:lax",
     ],
 )
 
diff --git a/jax/experimental/hijax.py b/jax/experimental/hijax.py
@@ -31,6 +31,9 @@
     HiPrimitive as HiPrimitive,
     HiType as HiType,
     MutableHiType as MutableHiType,
-    register_hitype as register_hitype,
     VJPHiPrimitive as VJPHiPrimitive,
+    register_hitype as register_hitype,
+)
+from jax._src.state import (
+    AbstractRef as AbstractRef,
 )
