Internal change

PiperOrigin-RevId: 673409076
jax-ml · Sep 11, 2024 · ea68f45 · ea68f45
1 parent 49dd6ed
commit ea68f45
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/tests/mosaic/BUILD b/tests/mosaic/BUILD
@@ -34,17 +34,23 @@ DISABLED_BACKENDS = [
 ]
 
 DISABLED_CONFIGS = [
-    "gpu",
     "gpu_a100",
+    "gpu_a100_x32",
     "gpu_p100",
     "gpu_p100_x32",
-    "gpu_x32",
     "gpu_pjrt_c_api",
+    "gpu_x32",
+    "gpu",
 ]
 
 jax_test(
     name = "gpu_test",
     srcs = ["gpu_test.py"],
+    config_tags_overrides = {
+        "gpu_h100_2gpu": {
+            "ondemand": False,  # Include in presubmit.
+        },
+    },
     disable_backends = DISABLED_BACKENDS,
     disable_configs = DISABLED_CONFIGS,
     enable_configs = ["gpu_h100_2gpu"],

diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
@@ -182,9 +182,8 @@ jax_test(
         "mosaic_gpu_test.py",
     ],
     config_tags_overrides = {
-        # TODO(slebedev): Switch to False once Mosaic GPU is unconditionally enabled.
         "gpu_h100_x32": {
-            "ondemand": True,  # Include in presubmit.
+            "ondemand": False,  # Include in presubmit.
         },
     },
     disable_backends = [