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

[jax2tf] Fix jax2tf tolerance for Cholesky, needed for newer TPUs

Browse Source 
PiperOrigin-RevId: 652751771
This commit is contained in:
George Necula 2024-07-16 02:12:27 -07:00 committed by jax authors
parent d34a6e9ce2
commit e7be205a39
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/experimental/jax2tf/tests/jax2tf_limitations.py
View File

@ -372,6 +372,13 @@ class Jax2TfLimitation(test_harnesses.Limitation):
tol=5e-2,
devices=("cpu", "gpu"),
modes=("eager", "graph", "compiled")),
custom_numeric(
dtypes=[dtypes.bfloat16],
tol=5e-5,
# Error for GL
devices=("tpu",),
modes=("eager", "graph", "compiled"),
native_serialization=Jax2TfLimitation.FOR_NATIVE),
custom_numeric(
custom_assert=custom_assert,
description=(