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

Merge pull request #23741 from awshaichen:neuron

Browse Source 
PiperOrigin-RevId: 678006682
This commit is contained in:
jax authors 2024-09-23 17:41:11 -07:00
commit 2ac1d0b8d0
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/compilation_cache.py
View File

@ -72,7 +72,7 @@ def is_cache_used(backend: xla_client.Client) -> bool:
# backend that supports serialization of executables.
# TODO(skye): add warning when initializing cache on unsupported default
# platform
supported_platforms = ["tpu", "gpu", "cpu"]
supported_platforms = ["tpu", "gpu", "cpu", "neuron"]
if not _is_cache_enabled():
monitoring.record_event('/jax/compilation_cache/task_disabled_cache')

2
jax/_src/interpreters/mlir.py
View File

@ -951,7 +951,7 @@ class LoweringResult(NamedTuple):
shape_poly_state: ShapePolyLoweringState
_platforms_with_donation = ["cpu", "cuda", "rocm", "tpu"]
_platforms_with_donation = ["cpu", "cuda", "rocm", "tpu", "neuron"]
def add_manual_axes(axis_ctx: sharding_impls.SPMDAxisContext, sharding, ndim):