Removed unnecessary BUILD dependencies from :ops_test

I also re-added the accidentally removed JAX_TRITON_COMPILE_VIA_XLA variable
to :pallas_test.
PiperOrigin-RevId: 621299158

Author: superbobry

jax/BUILD

@@ -104,13 +104,15 @@ package_group(
     name = "pallas_gpu_users",
     packages = [
         "//...",
+        "//learning/brain/research/jax",
     ] + pallas_gpu_internal_users,
 )
 
 package_group(
     name = "pallas_tpu_users",
     packages = [
         "//...",
+        "//learning/brain/research/jax",
     ] + pallas_tpu_internal_users,
 )
 

jax/_src/pallas/pallas_call.py

@@ -523,17 +523,23 @@ def _pallas_call_lowering(
   if platform == "cpu":
     raise ValueError("Only interpret mode is supported on CPU backend.")
   elif platform == "cuda" or platform == "rocm":
-    from jax._src.pallas.triton import pallas_call_registration  # type: ignore
-
-    return pallas_call_registration.pallas_call_lowering(
-        ctx, *in_nodes, interpret=interpret, **params
-    )
+    try:
+      from jax._src.pallas.triton import pallas_call_registration  # type: ignore
+    except ImportError:
+      pass
+    else:
+      return pallas_call_registration.pallas_call_lowering(
+          ctx, *in_nodes, interpret=interpret, **params
+      )
   elif platform == "tpu":
-    from jax._src.pallas.mosaic import pallas_call_registration  # type: ignore
-
-    return pallas_call_registration.pallas_call_tpu_lowering_rule(
-        ctx, *in_nodes, interpret=interpret, **params
-    )
+    try:
+      from jax._src.pallas.mosaic import pallas_call_registration  # type: ignore
+    except ImportError:
+      pass
+    else:
+      return pallas_call_registration.pallas_call_tpu_lowering_rule(
+          ctx, *in_nodes, interpret=interpret, **params
+      )
 
   raise _unsupported_lowering_error(platform)
 

jax/_src/pallas/triton/BUILD

@@ -67,5 +67,7 @@ pytype_strict_library(
         "//jax:util",
         "//jax/_src/lib",
         "//jax/_src/pallas",
-    ] + py_deps("jax_triton"),
+        # Users are expected to add a jax_triton dependency to use the legacy
+        # lowering path.
+    ],
 )

jax/_src/pallas/triton/pallas_call_registration.py

@@ -56,8 +56,8 @@ def compile_jaxpr(
     num_stages: int,
     debug: bool,
 ) -> CompilationResult:
-  from jax_triton.triton_lib import compile_ttir_to_ptx_inplace
-  import triton.backends.nvidia.compiler as cb
+  from jax_triton.triton_lib import compile_ttir_to_ptx_inplace  # type: ignore
+  import triton.backends.nvidia.compiler as cb  # type: ignore
 
   # TODO(sharadmv): handle multiple devices, right now we assume device 0
   # which is fine when we have multiple of the same GPU but this won't work in

jax/experimental/jax2tf/tests/primitives_test.py

@@ -182,6 +182,8 @@ def test_primitive_coverage(self):
       if p.name == "debug_callback":
         # TODO(sharadmv,necula): enable debug callbacks in TF
         continue
+      if p.name == "pallas_call":
+        continue
       if p.name in tf_not_yet_impl:
         self.assertNotIn(
             p, tf_impl)  # Should not be in both tf_impl and tf_not_yet_impl

tests/pallas/BUILD

@@ -56,10 +56,13 @@ jax_test(
         "gpu_a100_x32",
         "gpu_h100_x32",
     ],
+    env = {
+        "JAX_TRITON_COMPILE_VIA_XLA": "0",
+    },
     shard_count = 4,
     deps = [
         "//jax:pallas_gpu",
-    ] + py_deps("absl/testing") + py_deps("numpy"),
+    ] + py_deps("absl/testing") + py_deps("jax_triton") + py_deps("numpy"),
 )
 
 jax_test(
@@ -162,8 +165,7 @@ jax_test(
         "gpu_h100_x32",
     ],
     deps = [
-        "//jax:pallas_gpu",
-        "//jax:pallas_tpu",
+        "//jax:pallas",
     ] + py_deps("absl/testing") + py_deps("numpy"),
 )