mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26804 from hawkinsp:tsan

Browse Source 
PiperOrigin-RevId: 731721390
This commit is contained in:
jax authors 2025-02-27 07:39:35 -08:00
commit a8738a069e
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.github/workflows/tsan-suppressions.txt vendored
View File

@ -32,8 +32,6 @@ race:split_keys_entry_added
# https://github.com/python/cpython/issues/129748
race:mi_block_set_nextx
# https://github.com/python/cpython/issues/130571
race:_PyObject_GetMethod
# Races because the LAPACK and BLAS in our scipy isn't TSAN instrumented.
race:heevd_ffi
@ -63,3 +61,6 @@ race:gemm_oncopy
# https://github.com/python/cpython/issues/128133
# race:bytes_hash
# https://github.com/python/cpython/issues/130571
# race:_PyObject_GetMethod

2
tests/pjit_test.py
View File

@ -7495,10 +7495,12 @@ class UtilTest(jtu.JaxTestCase):
),
)
@jtu.thread_unsafe_test()
def test_op_sharding_cache_on_mesh_pspec_sharding(self):
ndim = 2
mesh = jtu.create_mesh((4, 2), ('x', 'y'))
mps1 = NamedSharding(mesh, P('x', 'y'))
sharding_impls.named_sharding_to_xla_hlo_sharding.cache_clear()
op1 = mps1._to_xla_hlo_sharding(ndim)
cache_info1 = sharding_impls.named_sharding_to_xla_hlo_sharding.cache_info()