From 169daec6b5ff4ce75213568e70d4a66935069273 Mon Sep 17 00:00:00 2001
From: Hana Joo <hanajoo@google.com>
Date: Fri, 22 May 2026 05:00:11 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 919610666
---
 BUILD              | 14 +++++++-------
 bazel/test_oss.bzl | 12 ++++++------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/BUILD b/BUILD
index 7ce857c..a34c56e 100644
--- a/BUILD
+++ b/BUILD
@@ -2,7 +2,7 @@
 
 load("@jaxite//bazel:test_oss.bzl", "cpu_gpu_tpu_test", "gpu_tpu_test", "multichip_tpu_test", "tpu_test")
 load("@rules_license//rules:license.bzl", "license")
-load("@rules_python//python:defs.bzl", "py_library", "py_test")
+load("@rules_python//python:defs.bzl", "py_library", "py_strict_test")
 
 package(
     default_applicable_licenses = [":license"],
@@ -38,7 +38,7 @@ py_library(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "encode_test",
     size = "small",
     timeout = "long",
@@ -53,7 +53,7 @@ py_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "encrypt_test",
     size = "small",
     timeout = "long",
@@ -430,7 +430,7 @@ cpu_gpu_tpu_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "lut_test",
     srcs = ["jaxite/jaxite_bool/lut_test.py"],
     deps = [
@@ -475,7 +475,7 @@ multichip_tpu_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "rns_test",
     size = "small",
     timeout = "moderate",
@@ -493,7 +493,7 @@ py_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "rns_utils_test",
     size = "small",
     timeout = "moderate",
@@ -570,7 +570,7 @@ tpu_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "math_test",
     size = "small",
     timeout = "moderate",
diff --git a/bazel/test_oss.bzl b/bazel/test_oss.bzl
index 918ae70..b42acaa 100644
--- a/bazel/test_oss.bzl
+++ b/bazel/test_oss.bzl
@@ -1,6 +1,6 @@
 """Testing util helpers"""
 
-load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:defs.bzl", "py_strict_test")
 
 # This file is a shim to have the same Google-internal and external API
 # for running jaxite on TPU and GPU tests. The internal API is not
@@ -25,7 +25,7 @@ def tpu_test(name, **kwargs):
     if not kwargs.get("main"):
         main = kwargs.get("srcs", ["%s.py" % name])[0]
         kwargs.setdefault("main", main)
-    py_test(
+    py_strict_test(
         name = name,
         srcs = kwargs["srcs"],
         size = kwargs["size"],
@@ -51,7 +51,7 @@ def cpu_tpu_test(name, **kwargs):
       name: the name of the CPU test.
       **kwargs: other args passed along to generated rules.
     """
-    py_test(name = name, **kwargs)
+    py_strict_test(name = name, **kwargs)
 
 def gpu_tpu_test(name, **kwargs):
     """A shim that only generates CPU tests.
@@ -60,7 +60,7 @@ def gpu_tpu_test(name, **kwargs):
       name: the name of the CPU test.
       **kwargs: other args passed along to generated rules.
     """
-    py_test(name = name, **kwargs)
+    py_strict_test(name = name, **kwargs)
 
 def cpu_gpu_tpu_test(name, **kwargs):
     """A shim that only generates CPU tests.
@@ -69,7 +69,7 @@ def cpu_gpu_tpu_test(name, **kwargs):
       name: the name of the CPU test.
       **kwargs: other args passed along to generated rules.
     """
-    py_test(name = name, **kwargs)
+    py_strict_test(name = name, **kwargs)
 
 def cpu_gpu_tpu_py_benchmark(name, **kwargs):
     """A shim that only generates CPU tests.
@@ -79,4 +79,4 @@ def cpu_gpu_tpu_py_benchmark(name, **kwargs):
         `_<platform>_benchmark`, e.g., `_cpu_benchmark`
       **kwargs: other args passed along to generated rules.
     """
-    py_test(name = name + "_cpu_benchmark", **kwargs)
+    py_strict_test(name = name + "_cpu_benchmark", **kwargs)