jax
35d75183 - `_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.

Commit

304 days ago

`_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation. PiperOrigin-RevId: 743986126

References

#27668 - `_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.

Author

Google-ML-Automation

Committer

Google-ML-Automation

Parents

b9007145

jax 35d75183 - `_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.

jax
35d75183 - `_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.