Merge pull request #25097 from jburnim:jburnim_pallas_interpret_mode

PiperOrigin-RevId: 724073443
jax-ml · Feb 6, 2025 · 4b86ff2 · 4b86ff2
2 parents 840192d + 1c82484
commit 4b86ff2
Show file tree

Hide file tree

Showing 7 changed files with 2,040 additions and 6 deletions.
diff --git a/jax/BUILD b/jax/BUILD
@@ -652,6 +652,7 @@ pytype_strict_library(
         "//jax/_src/pallas",
         "//jax/_src/pallas/mosaic:core",
         "//jax/_src/pallas/mosaic:helpers",
+        "//jax/_src/pallas/mosaic:interpret",
         "//jax/_src/pallas/mosaic:lowering",
         "//jax/_src/pallas/mosaic:pallas_call_registration",  # build_cleaner: keep
         "//jax/_src/pallas/mosaic:pipeline",

diff --git a/jax/_src/pallas/mosaic/BUILD b/jax/_src/pallas/mosaic/BUILD
@@ -148,3 +148,15 @@ py_library(
         "//jax/_src/pallas",
     ],
 )
+
+py_library(
+    name = "interpret",
+    srcs = ["interpret.py"],
+    deps = [
+        ":core",
+        ":primitives",
+        "//jax",
+        "//jax/_src/lib",
+        "//jax/_src/pallas",
+    ] + py_deps("numpy"),
+)