[benikm91]

API change: Before → now

Operation	JAX	Before	Now
Index	x[i]	t.slice(Axis[L].at(i))	unchanged (removes L)
Traced index	x[i], lax.dynamic_index_in_dim	t.slice(Axis[L].at(i: Tensor0[Int32]))	unchanged (removes L)
Range	x[a:b:s], lax.slice	t.slice(Axis[L].at(a until b by s))	unchanged (keeps L)
List of indices	x[[i, j]]	t.slice(Axis[L].at(Seq(i, j))) or at((i, j))	unchanged (keeps L)
Vector of indices	jnp.take(x, idx, axis)	t.take(Axis[L])(idx)	t.slice(Axis[L].at(idx: Tensor1[L2, Int32])), L replaced by L2
Window	lax.dynamic_slice	Tensor1 only: t.dynamicSlice(start, size)	t.slice(Axis[L].at(start, windowSize)), any rank and axis; start is an Int or a Tensor0[Int32]
Several axes	x[i, a:b]	t.slice((Axis[L].at(i), Axis[K].at(a until b)))	removed: t.slice(Axis[L].at(i)).slice(Axis[K].at(a until b))
Set	x.at[...].set(v)	t.set(selector)(v) for index, traced index, range, list, or a tuple	t.set(selector)(v) for every selector above, including a vector of indices, at((i, j)) and a window (lax.dynamic_update_slice)
Set a scalar	x.at[i, j].set(5.0)	t.set((Axis[A].at(i), Axis[B].at(j)))(5f)	removed:

Breaking changes

Remove tuple selectors from slice and set

• It had a bug. All selectors were combined into one NumPy-style index, so two index lists on different axes were paired element by element. For example, t.slice((Axis[A].at(Seq(0, 1)), Axis[B].at(Seq(0, 1)))) returned the diagonal: runtime shape (2,), while the type claimed (A, B). Chained calls select every combination, as the type says.
• It made the implementation unnecessarily complex. The new window lowers to lax.dynamic_slice, not to indexing, so it could only have joined the tuple through special cases in the index builder.
• Nothing used it outside DimWit's own tests: not the examples, not DeepWit.

Other breaking changes

• t.take(Axis[L])(idx) → t.slice(Axis[L].at(idx))
• t.dynamicSlice(start, size) (Tensor1 only) → t.slice(Axis[L].at(start, size))
• Axis[L].at(a, b) without double parentheses now selects a window; a tuple of indices is written Axis[L].at((a, b)).
• set requires the value to have the shape of the selection, instead of letting JAX broadcast it.
• AxesConditionalRemover is removed; only the tuple form used it.

[marcelluethi]

I like the idea of cleaning up the selectors. But I am not sure if using slice for everything is a move in the right direction. Slice has now many different use cases, from selecting a single element over a continuous region to selecting irregularly spaced elements in the tensor.

I wonder if we could instead use the following convention:

• slice: for a regular region (index, range with step, window with static or traced start)
• take: for selection by an index list or index tensor
• at/get: a single scalar, given one coordinate per axis; returns a Tensor0

To me the difference of take/slice seems meaningful, as I would not associate taking elements 1, 3, 7 from a 1d-tensor with slice.

[benikm91]

I agree that slice is a bad name for the take case. I'd still keep one method, because an index, a range or an index tensor only change which axes the result has, not what the operation does. select fits every case. Adding apply on top lets us write grid(Axis[A].at(1)), which mirrors Scala collections, e.g. array(i).

Additionally, set sounds like mutation, but the operation returns a new tensor. Scala's immutable collections, such as Vector and Map, use: apply to read and updated to write.

So, I propose:

• Keep this PR changes but rename slice to select.
• Add apply, forwarding to select.
• Rename set to updated

[benikm91]

I agree that slice is a bad name for the take case. I'd still keep one method, because an index, a range or an index tensor only change which axes the result has, not what the operation does. select fits every case. Adding apply on top lets us write grid(Axis[A].at(1)), which mirrors Scala collections, e.g. array(i).

Additionally, set sounds like mutation, but the operation returns a new tensor. Scala's immutable collections, such as Vector and Map, use: apply to read and updated to write.

So, I propose:

• Keep this PR changes but rename slice to select.
• Add apply, forwarding to select.
• Rename set to updated

@marcelluethi I added a commit with my proposal

[marcelluethi]

This looks good to me. Maybe we should keep the old methods (slice and take) around, forward their implementation to select and mark them as deprecated? what do you think?

[marcelluethi]

Here one slice is replaced by two selects. Is this intended and do we know how it affects the performance / memory overhead?

[benikm91]

Yes this was dropped in the first commit of this PR. I removed it as the main branch version is partly buggy and led to very complex implementation of "slice" - and this was never used so far (See "Remove tuple selectors from slice and set" in first comment of PR). It could be that doing two slices is slower. I also removed the same for set where there is no easy alternative.

The motivation of removal is that I overdid it in the first implementation (current main), leading to a partly correct, badly tested, and badly implemented, yet general version. We will probably add this generality back at some point, but I see it less important right now than focusing on keeping the supported features correct and clean.

[benikm91]

@marcelluethi As this comment highlights a breaking change in this PR I requested a re-review. To check if you are fully on board with this before I merge into main. I did not change code.

[benikm91]

This looks good to me. Maybe we should keep the old methods (slice and take) around, forward their implementation to select and mark them as deprecated? what do you think?

No we should not. I am for removing over deprecation as long as we are below v1.0.0.