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

Fix test failures.

Browse Source 
PiperOrigin-RevId: 662703221
This commit is contained in:
Peter Hawkins 2024-08-13 17:01:33 -07:00 committed by jax authors
parent 2b3ccce793
commit 323e257f67
3 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/lax_numpy_reducers_test.py
View File

@ -715,7 +715,7 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
# TODO(phawkins): we currently set dtype=False because we aren't as
# aggressive about promoting to float64. It's not clear we want to mimic
# Numpy here.
tol_spec = {np.float16: 1E-2, np.float32: 2e-4, np.float64: 5e-6}
tol_spec = {np.float16: 4e-2, np.float32: 2e-4, np.float64: 5e-6}
tol = max(jtu.tolerance(a_dtype, tol_spec),
jtu.tolerance(q_dtype, tol_spec))
self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,

2
tests/lax_numpy_test.py
View File

@ -4768,7 +4768,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
else x.astype(np.float32) for x in choicelist]
dtype = jnp.result_type(default, *choicelist)
return np.select(condlist,
[np.asarray(x, dtype=dtype) for x in choicelist],
[np.asarray(x).astype(dtype) for x in choicelist],
np.asarray(default, dtype=dtype))
with jtu.strict_promotion_if_dtypes_match(dtypes):
self._CheckAgainstNumpy(np_fun, jnp.select, args_maker,

3
tests/lax_scipy_test.py
View File

@ -226,7 +226,8 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
rng = jtu.rand_positive(self.rng())
args_maker = lambda: [rng(shape, dtype) + (d - 1) / 2.]
self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker,
tol={np.float32: 1e-3, np.float64: 1e-14})
tol={np.float32: 1e-3, np.float64: 1e-14},
check_dtypes=False)
self._CompileAndCheck(
lax_fun, args_maker, rtol={
np.float32: 5e-5 if jtu.test_device_matches(["tpu"]) else 1e-05,