mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix type annotation for bcoo_update_layout.

Browse Source 
PiperOrigin-RevId: 493567424
This commit is contained in:
Shaobo Hou 2022-12-07 04:24:29 -08:00 committed by jax authors
parent ac72346ad3
commit 6e0c8029cc
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/experimental/sparse/bcoo.py
View File

@ -1565,7 +1565,7 @@ mlir.register_lowering(bcoo_sum_duplicates_p, _bcoo_sum_duplicates_mhlo)
# BCOO functions that maybe should be primitives?
def bcoo_update_layout(mat: BCOO, *, n_batch: Optional[int] = None, n_dense: Optional[int] = None,
on_inefficient: str = 'error') -> BCOO:
on_inefficient: Optional[str] = 'error') -> BCOO:
"""Update the storage layout (i.e. n_batch & n_dense) of a BCOO matrix.
In many cases this can be done without introducing undue storage overhead. However,