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

Slice.from_slice now works for slices with a negative start index

Browse Source 
The implementation still requires the step to be 1, so any slice
with a negative start index has size 0.
This commit is contained in:
Sergei Lebedev 2024-01-29 12:18:18 +00:00
parent f910a0d4f8
commit 3cea57d9d1
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/_src/state/indexing.py
View File

@ -53,7 +53,7 @@ class Slice:
start, stop, step = slc.indices(size)
if step != 1:
raise ValueError(f"slice must have a step of 1 (found: {step})")
return cls(start, stop - start)
return cls(start, max(stop - start, 0))
def dslice(start: int | Array | None, size: int | None = None