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

Add instructions to report issues to GitHub for unsupported features in PJRT C API.

Browse Source 
PiperOrigin-RevId: 629572930
This commit is contained in:
Jieying Luo 2024-04-30 17:14:49 -07:00 committed by jax authors
parent 738612802e
commit 9b79141ef2
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/experimental/host_callback.py
View File

@ -586,9 +586,8 @@ def _raise_if_using_outfeed_with_pjrt_c_api(backend: xb.XlaBackend):
"""Should be called whenever outfeed (or infeed) will be used."""
if xb.using_pjrt_c_api(backend):
raise NotImplementedError(
"host_callback functionality isn't supported with the new Cloud TPU "
"runtime. See https://jax.readthedocs.io/en/latest/debugging/index.html"
" and "
"host_callback functionality isn't supported with PJRT C API. "
"See https://jax.readthedocs.io/en/latest/debugging/index.html and "
"https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html"
" for alternatives. Please file a feature request at "
"https://github.com/google/jax/issues if none of the alternatives are "