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

[ROCm] Add rocm support in nm.py lowering code

Browse Source
This commit is contained in:
Ruturaj4 2024-06-05 20:40:35 +00:00
parent 9e1a4e8d74
commit d2ab42de3d
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/experimental/sparse/nm.py
View File

@ -181,6 +181,9 @@ dispatch.simple_impl(nm_spmm_p)
if gpu_sparse.cuda_is_supported:
mlir.register_lowering(nm_spmm_p, _nm_spmm_gpu_lowering, platform="cuda")
if gpu_sparse.rocm_is_supported:
mlir.register_lowering(nm_spmm_p, _nm_spmm_gpu_lowering, platform="ROCM")
# --------------------------------------------------------------------
# nm_pack