jax.lax.dynamic_slice accepting numpy slice syntax

[jglaser]

The API for jax.lax.dynamic_slice() requires tuples of start and length indices... I do not understand the reason why the API is different from the usual numpy slice syntax, but I admittedly also do not understand why there is a different API for dynamic slicing at all.

Perhaps an API and/or a documentation update would be useful?

[jakevdp]

Hi - the reason this is different is because in dynamic_slice, the start index is dynamic, but the slice size is static. There is no way to construct such a slice using NumPy's slice syntax; for example, if you have a dynamic start_index and a static length, then you would have to write something like this:

end_index = start_index + length
x[start_index:end_index]

Because end_index is derived from start_index, it is dynamic, and your indexing expression no longer has any way to statically determine the length!

We may someday introduce a higher-level syntax for this: for example pallas has a dynamic slice object that makes dynamic-slice indexing look something like x[dslice(start_index, length)], but it's not clear whether we'll adopt that kind of helper in JAX itself. In the meantime, if you want to dynamically slice an array, you can call the dynamic_slice API directly:

jax.lax.dynamic_slice(x, (start_index,), (length),)

[davisyoshida]

[jglaser]

Thanks for the pointer, looks interesting, although my use case is a little different --- linear algebra on sharded matrices, implementing optimized, blocked versions of Cholesky decomposition and triangular solves. Those algorithms require you to broadcast diagonal blocks to other processors in every loop iteration, and hence the dynamic index - if the loop over diagonal blocks is implemented using jax.lax.fori_loop to save compile time

[davisyoshida]

Actually I had meant to post this to another discussion, sorry to waste your time!

[patrick-kidger]

Since you mention doing linear solvers -- you might find Lineax interesting, which we're now using as our standard API abstraction in downstream scientific libraries (Diffrax, Optimistix, etc.)

We haven't yet really focused on sharded/blocked approaches though! So I'd be curious to see if there's something more we can do on that front. :)