Merge pull request #26185 from hawkinsp:tsan

PiperOrigin-RevId: 721063414

.github/workflows/tsan-suppressions.txt

@@ -11,28 +11,51 @@ race:ensure_nonmanaged_dict
 # https://github.com/openxla/xla/issues/20686
 race:dnnl_sgemm
 
-# https://github.com/numpy/numpy/issues/28041
-race:get_initial_from_ufunc
-
-# https://github.com/numpy/numpy/issues/28042
-race:PyArray_UpdateFlags
-
 # https://github.com/python/cpython/issues/128130
 race_top:run_eval_code_obj
 
+# Likely only happens when the process is crashing.
 race:dump_traceback
 
-# https://github.com/numpy/numpy/issues/28045  not sure about this one
-race:arraymethod_dealloc
-
-# https://github.com/python/cpython/issues/128133
-race:bytes_hash
-
 # https://github.com/python/cpython/issues/128137
+# Fixed in Python 3.14, but not backported to 3.13.
 race:immortalize_interned
 
 # https://github.com/python/cpython/issues/128144
+# Fixed in Python 3.14, but not backported to 3.13.
 race_top:PyMember_GetOne
 
+# https://github.com/python/cpython/issues/128133
+# Fixed in Python 3.14, but not backported to 3.13.
+race:bytes_hash
+
 # https://github.com/python/cpython/issues/128657
 race:py_digest_by_name
+
+# https://github.com/python/cpython/issues/128714
+race:func_get_annotations
+
+# Races because the LAPACK and BLAS in our scipy isn't TSAN instrumented.
+race:heevd_ffi
+race:gesdd_ffi
+race:dscal_k_
+race:scal_k_
+race:gemm_beta
+race:gemm_oncopy
+
+
+# Races that are probably fixed
+# https://github.com/numpy/numpy/issues/28041
+# race:get_initial_from_ufunc
+#
+# https://github.com/numpy/numpy/issues/28042
+# race:PyArray_UpdateFlags
+#
+# https://github.com/python/cpython/issues/128717
+# race:Py_SetRecursionLimit
+#
+# https://github.com/numpy/numpy/issues/28143
+# race:PyArray_CheckFromAny_int
+
+# https://github.com/python/cpython/issues/128759
+# race:PyType_Modified

.github/workflows/tsan.yaml

@@ -9,8 +9,8 @@ on:
     - cron: "0 12 * * *" # Daily at 12:00 UTC
   workflow_dispatch: # allows triggering the workflow run manually
   pull_request: # Automatically trigger on pull requests affecting this file
-    # branches:
-    #   - main
+    branches:
+      - main
     paths:
       - '**/workflows/tsan.yaml'
 
@@ -99,4 +99,6 @@ jobs:
               --test_env=TSAN_OPTIONS=halt_on_error=1,suppressions=$PWD/.github/workflows/tsan-suppressions.txt \
               --test_env=JAX_TEST_NUM_THREADS=8 \
               --test_output=errors \
+              --local_test_jobs=32 \
+              --test_timeout=600 \
               //tests:cpu_tests