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

Test: fix casting warning in betainc test

Browse Source
This commit is contained in:
Jake VanderPlas 2023-12-04 14:20:14 -08:00
parent c4d2fc7364
commit 8b74b93501
2 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/_src/internal_test_util/lax_test_util.py
View File

@ -278,8 +278,11 @@ def lax_ops():
"betainc",
3,
float_dtypes,
test_util.rand_positive,
{np.float64: 1e-14},
test_util.rand_uniform,
{
np.float32: 1e-5,
np.float64: 1e-12,
},
),
op_record(
"igamma",

2
tests/lax_test.py
View File

@ -105,7 +105,7 @@ class LaxTest(jtu.JaxTestCase):
tol = tol or jtu.default_tolerance()
if jtu.test_device_matches(["tpu"]):
if dtype in (np.float32, np.complex64) and op_name in (
"acosh", "asinh", "cos", "cosh", "digamma", "exp", "exp2", "igamma",
"acosh", "asinh", "betainc", "cos", "cosh", "digamma", "exp", "exp2", "igamma",
"igammac", "log", "log1p", "logistic", "pow", "sin", "sinh", "tan"):
tol = jtu.join_tolerance(tol, 2e-4)
elif op_name == "asinh" and dtype == np.float16: