From ce639ea9b5285c0f8cad0c7759e1fe351dbcbafe Mon Sep 17 00:00:00 2001
From: Hana Joo <hanajoo@google.com>
Date: Thu, 7 May 2026 05:06:17 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 911900930
---
 BUILD              | 13 +++++++------
 bazel/test_oss.bzl |  2 +-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/BUILD b/BUILD
index 7ce857c..a40e367 100644
--- a/BUILD
+++ b/BUILD
@@ -2,7 +2,8 @@
 
 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")
+load("//devtools/python/blaze:strict.bzl", "py_strict_test")
 
 package(
     default_applicable_licenses = [":license"],
@@ -38,7 +39,7 @@ py_library(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "encode_test",
     size = "small",
     timeout = "long",
@@ -430,7 +431,7 @@ cpu_gpu_tpu_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "lut_test",
     srcs = ["jaxite/jaxite_bool/lut_test.py"],
     deps = [
@@ -475,7 +476,7 @@ multichip_tpu_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "rns_test",
     size = "small",
     timeout = "moderate",
@@ -493,7 +494,7 @@ py_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "rns_utils_test",
     size = "small",
     timeout = "moderate",
@@ -570,7 +571,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..f66d90d 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("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_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