[Pallas/interpreter] Refactor re-useable functionality out of the main source file for the TPU kernel interpreter.

PiperOrigin-RevId: 845731033

Author: Google-ML-Automation

jax/_src/pallas/mosaic/interpret/BUILD

@@ -33,6 +33,8 @@ py_library(
     deps = [
         ":race_detection_state",
         ":shared_memory",
+        ":thread_map",
+        ":utils",
         ":vector_clock",
         "//jax",
         "//jax/_src:api",
@@ -79,3 +81,23 @@ pytype_strict_library(
         "//jax/_src:source_info_util",
     ],
 )
+
+pytype_strict_library(
+    name = "thread_map",
+    srcs = ["thread_map.py"],
+    deps = [
+        "//jax",
+        "//jax/_src:callback",
+    ],
+)
+
+pytype_strict_library(
+    name = "utils",
+    srcs = ["utils.py"],
+    deps = [
+        "//jax",
+        "//jax/_src:core",
+        "//jax/_src:util",
+        "//jax/_src/pallas",
+    ] + py_deps("numpy"),
+)