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

update sparse _reshape docstring to match jax.Array.reshape

Browse Source
This commit is contained in:
Blair-Johnson 2024-02-08 16:15:00 -05:00
parent 4c505f8bac
commit 8bcd3d1049
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/transform.py
View File

@ -889,7 +889,7 @@ def _sum(self, *args, **kwargs):
return sparsify(lambda x: x.sum(*args, **kwargs))(self)
def _reshape(self, *args, **kwargs):
"""Sum array along axis."""
"""Returns an array containing the same data with a new shape."""
return sparsify(lambda x: x.reshape(*args, **kwargs))(self)
def _astype(self, *args, **kwargs):