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

Reenable GPU int matmul test since the XLA bug is fixed.

Browse Source 
PiperOrigin-RevId: 495117439
This commit is contained in:
Peter Hawkins 2022-12-13 13:50:10 -08:00 committed by jax authors
parent c6eb632f57
commit b80af85298
1 changed files with 0 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/lax_vmap_test.py
View File

@ -219,8 +219,6 @@ class LaxVmapTest(jtu.JaxTestCase):
dtype=default_dtypes,
)
def testDot(self, lhs_shape, rhs_shape, dtype, bdims):
if jtu.device_under_test() == "gpu" and dtype == np.int64:
raise unittest.SkipTest("Wrong outputs for batched matmuls (b/258497059)")
rng = jtu.rand_default(self.rng())
op = partial(lax.dot, precision=lax.Precision.HIGHEST)
self._CheckBatching(op, 5, bdims, (lhs_shape, rhs_shape), (dtype, dtype),