mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 23:46:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update scatter.py

Browse Source 
Looks like the documentation was not up to date
This commit is contained in:
AdrienCorenflos 2020-11-05 18:13:44 +02:00 committed by GitHub
parent 11d6f239fd
commit 3bc34e4f5a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
jax/_src/ops/scatter.py
View File

@ -207,8 +207,8 @@ def index_min(x, idx, y, indices_are_sorted=False, unique_indices=False):
:data:`jax.ops.index` object.
y: the array of updates. `y` must be broadcastable to the shape of the
array that would be returned by `x[idx]`.
indices_are_sorted: whether `scatter_indices` is known to be sorted
unique_indices: whether `scatter_indices` is known to be free of duplicates
indices_are_sorted: whether `idx` is known to be sorted
unique_indices: whether `idx` is known to be free of duplicates
Returns:
An array.
@ -246,8 +246,8 @@ def index_max(x, idx, y, indices_are_sorted=False, unique_indices=False):
:data:`jax.ops.index` object.
y: the array of updates. `y` must be broadcastable to the shape of the
array that would be returned by `x[idx]`.
indices_are_sorted: whether `scatter_indices` is known to be sorted
unique_indices: whether `scatter_indices` is known to be free of duplicates
indices_are_sorted: whether `idx` is known to be sorted
unique_indices: whether `idx` is known to be free of duplicates
Returns:
An array.
@ -286,8 +286,8 @@ def index_update(x, idx, y, indices_are_sorted=False, unique_indices=False):
:data:`jax.ops.index` object.
y: the array of updates. `y` must be broadcastable to the shape of the
array that would be returned by `x[idx]`.
indices_are_sorted: whether `scatter_indices` is known to be sorted
unique_indices: whether `scatter_indices` is known to be free of duplicates
indices_are_sorted: whether `idx` is known to be sorted
unique_indices: whether `idx` is known to be free of duplicates
Returns:
An array.