[Pallas/interpreter] Move SharedMemory out of the main interpreter source file.

jax/_src/pallas/core.py

@@ -1443,7 +1443,7 @@ def _core_map_abstract_eval(*args, jaxpr, mesh, **kwargs):
   effs = set()
   if interpret:
     try:
-      from jax._src.pallas.mosaic import interpret as mosaic_tpu_interpret  # Avoid circular dependency.
+      from jax._src.pallas.mosaic.interpret import interpret_pallas_call as mosaic_tpu_interpret  # Avoid circular dependency.
       if isinstance(interpret, mosaic_tpu_interpret.InterpretParams):
         effs = mosaic_tpu_interpret.get_interpret_effects()
     except ImportError:
@@ -1619,7 +1619,7 @@ def _core_map_typecheck_rule(_, *in_atoms, jaxpr, mesh, **kwargs):
   effs = set()
   if interpret:
     try:
-      from jax._src.pallas.mosaic import interpret as mosaic_tpu_interpret  # Avoid circular dependency.
+      from jax._src.pallas.mosaic.interpret import interpret_pallas_call as mosaic_tpu_interpret  # Avoid circular dependency.
       if isinstance(interpret, mosaic_tpu_interpret.InterpretParams):
         effs = mosaic_tpu_interpret.get_interpret_effects()
     except ImportError:

jax/_src/pallas/mosaic/BUILD

@@ -230,19 +230,3 @@ py_library(
         "//jax/_src/pallas",
     ],
 )
-
-py_library(
-    name = "interpret",
-    srcs = ["interpret.py"],
-    deps = [
-        ":core",
-        ":primitives",
-        ":verification",
-        "//jax",
-        "//jax/_src:core",
-        "//jax/_src:source_info_util",
-        "//jax/_src:util",
-        "//jax/_src/lib",
-        "//jax/_src/pallas",
-    ] + py_deps("numpy"),
-)

jax/_src/pallas/mosaic/interpret/BUILD

@@ -0,0 +1,79 @@
+# Copyright 2025 The JAX Authors.
+#
+# 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.
+
+# Package for Pallas TPU Interpret Mode
+
+load("@rules_python//python:defs.bzl", "py_library")
+load("//jaxlib:jax.bzl", "py_deps", "pytype_strict_library")
+
+package(
+    default_applicable_licenses = [],
+    default_visibility = [
+        "//jax:internal",
+    ],
+)
+
+py_library(
+    name = "interpret_pallas_call",
+    srcs = ["interpret_pallas_call.py"],
+    deps = [
+        ":race_detection_state",
+        ":shared_memory",
+        ":vector_clock",
+        "//jax",
+        "//jax/_src:api",
+        "//jax/_src:callback",
+        "//jax/_src:config",
+        "//jax/_src:core",
+        "//jax/_src:frozen_dict",
+        "//jax/_src:lax",
+        "//jax/_src:mlir",
+        "//jax/_src:source_info_util",
+        "//jax/_src:typing",
+        "//jax/_src:util",
+        "//jax/_src/pallas",
+        "//jax/_src/pallas/mosaic:core",
+        "//jax/_src/pallas/mosaic:primitives",
+        "//jax/_src/pallas/mosaic:verification",
+    ] + py_deps("numpy"),
+)
+
+pytype_strict_library(
+    name = "vector_clock",
+    srcs = ["vector_clock.py"],
+    deps = py_deps("numpy"),
+)
+
+pytype_strict_library(
+    name = "shared_memory",
+    srcs = ["shared_memory.py"],
+    deps = [
+        ":race_detection_state",
+        ":vector_clock",
+        "//jax",
+        "//jax/_src:source_info_util",
+        "//jax/_src:typing",
+        "//jax/_src/pallas",
+        "//jax/_src/pallas/mosaic:core",
+    ] + py_deps("numpy"),
+)
+
+pytype_strict_library(
+    name = "race_detection_state",
+    srcs = ["race_detection_state.py"],
+    deps = [
+        ":vector_clock",
+        "//jax/_src:source_info_util",
+    ],
+)

jax/_src/pallas/mosaic/interpret/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 The JAX Authors.
+#
+# 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.