Unify slicing on single-axis Axis[L].at selectors: remove the tuple f…

AGENTS.md

@@ -647,21 +647,115 @@ println(s"Transposed shape: ${transposed.shape}")
 val vec1d = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f))
 val vec2d = vec1d.appendAxis(Axis[B])  // Add new axis B
 
-// Take (indexing with single index)
-val firstRow = original.slice(Axis[A].at(0))
-val secondRow = original.slice(Axis[A].at(1))
+// Select at an index (removes the axis)
+val firstRow = original.select(Axis[A].at(0))
+val secondRow = original.select(Axis[A].at(1))
 
-// Take (indexing with range)
+// Select at a range (keeps the axis)
 val data3d = Tensor(Shape3(Axis[A] -> 5, Axis[B] -> 3, Axis[C] -> 4)).fill(1.0f)
-val middleSlice = data3d.slice(Axis[A].at(1 until 4))  // Takes indices 1, 2, 3
-println(s"Middle slice shape: ${middleSlice.shape}")  // Shape(A -> 3, B -> 3, C -> 4)
+val middleSelection = data3d.select(Axis[A].at(1 until 4))  // Takes indices 1, 2, 3
+println(s"Middle selection shape: ${middleSelection.shape}")  // Shape(A -> 3, B -> 3, C -> 4)
 
 // Concatenate
 val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f))
 val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f))
 val concatenated = concatenate(t1, t2, Axis[A])
 ```
 
+### Selecting and Updating
+
+Every selection along an axis is an `Axis[L].at(...)` selector. `select` reads the selection, and applying
+the tensor to the selector, `t(selector)`, is the same as `t.select(selector)`, as `seq(i)` reads a Scala collection.
+`updated` returns a copy with new values at the selection, as `updated` does for Scala's immutable collections.
+A selector selects along exactly one axis; to select along several axes, chain the calls.
+
+```scala
+// Shape: A=4, B=3, the value at (a, b) is 10a + b
+val grid = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(4, 3)((a, b) => 10.0f * a + b))
+
+// Index: removes the axis
+val gridRow: Tensor1[B, Float32] = grid.select(Axis[A].at(1)) // [10, 11, 12]
+
+// Range, sequence or tuple of indices: keeps the axis, with the number of indices as extent
+val gridRows = grid.select(Axis[A].at(1 until 3)) // rows 1, 2
+val gridPicked = grid.select(Axis[A].at(Seq(3, 0))) // rows 3, 0
+val gridTupled = grid.select(Axis[A].at((3, 0))) // the same, note the double parentheses
+
+// Vector of indices, like jnp.take: the axis is replaced by the axis of the indices
+val order = Tensor1(Axis[C]).fromArray(Array(2, 0, 2))
+val gridGathered: Tensor2[A, C, Float32] = grid.select(Axis[B].at(order))
+
+// Window of static size, like jax.lax.dynamic_slice: keeps the axis with the window size as extent
+val gridWindow = grid.select(Axis[A].at(1, 2)) // rows 1, 2
+
+// Several axes: chain the calls
+val gridEntry: Tensor0[Float32] = grid.select(Axis[A].at(1)).select(Axis[B].at(2)) // 12
+
+// Apply: the same as select
+val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 12
+
+// updated takes a value with the shape of the selection
+val gridZeroRow = grid.updated(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f)))
+val gridZeroWindow = grid.updated(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f))
+```
+
+The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one,
+e.g. a loop position under `jit`. Its value is then only known at runtime, but the window size stays static:
+
+```scala
+// Write a row of ones into a buffer at a position known only at runtime
+val writeRow = jit((buffer: Tensor2[A, B, Float32], position: Tensor0[Int32]) =>
+  buffer.updated(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f))
+)
+val gridWritten = writeRow(grid, Tensor0(2)) // row 2 is now ones
+```
+
+```scala
+// ERROR: a selector selects along one axis, chain the calls instead
+val gridTuple = grid.select((Axis[A].at(1), Axis[B].at(2)))
+// error:
+// value select is not a member of dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B,
+//   dimwit.tensor.DType.Float32].
+// An extension method was tried, but could not be fully constructed:
+// 
+//     dimwit.select()
+// 
+//     failed with:
+// 
+//         value select: <overloaded dimwit.select> does not take parameters
+// val gridTuple = grid.select((Axis[A].at(1), Axis[B].at(2)))
+//                 ^^^^^^^^^^^
+```
+
+```scala
+// ERROR: an index removes the axis, so the value must not have it
+val gridWrongValue = grid.updated(Axis[A].at(0))(grid)
+// error:
+// Found:    (MdocApp1.this.grid :
+//   dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B,
+//     dimwit.tensor.DType.Float32]
+// )
+// Required: dimwit.tensor.Tensor[?1.RemainingAxes, dimwit.tensor.DType.Float32]
+// 
+// where:    ?1 is an unknown value of type dimwit.tensor.ShapeTypeHelpers.AxisRemover.Aux[
+//   (MdocApp1.this.A, MdocApp1.this.B), MdocApp1.this.A, MdocApp1.this.B *:
+//   EmptyTuple.type]
+// 
+// val gridWrongValue = grid.updated(Axis[A].at(0))(grid)
+//                                                  ^^^^
+```
+
+```scala
+// ERROR at runtime: a window with an Int start must lie inside the axis
+grid.select(Axis[A].at(3, 2))
+// java.lang.IllegalArgumentException: requirement failed: Window 3 until 5 is out of bounds for axis A of size 4
+// 	at scala.Predef$.require(Predef.scala:389)
+// 	at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:742)
+// 	at dimwit.tensor.tensorops.StructuralExtensions$.select(StructuralOps.scala:605)
+// 	at dimwit.package$.select(package.scala:80)
+// 	at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:723)
+```
+
 ---
 
 ## Functional Patterns

core/src/main/scala/dimwit/package.scala

@@ -4,6 +4,8 @@ import dimwit.tensor.AxisAtIndex
 import dimwit.tensor.AxisAtIndices
 import dimwit.tensor.AxisAtRange
 import dimwit.tensor.AxisAtTensorIndex
+import dimwit.tensor.AxisAtTensorIndices
+import dimwit.tensor.AxisAtWindow
 import dimwit.tensor.AxisExtent
 import dimwit.tensor.AxisSelector
 
@@ -55,9 +57,11 @@ package object dimwit:
     AxisAtIndex,
     AxisAtRange,
     AxisAtIndices,
-    AxisAtTensorIndex
+    AxisAtTensorIndex,
+    AxisAtTensorIndices,
+    AxisAtWindow
   }
-  export dimwit.tensor.ShapeTypeHelpers.{AxisInTensor, AxisIndex, AxisRemover, AxisReplacer, AxisIndices, AxesRemover, AxesConditionalRemover, SharedAxisRemover}
+  export dimwit.tensor.ShapeTypeHelpers.{AxisInTensor, AxisIndex, AxisRemover, AxisReplacer, AxisIndices, AxesRemover, SharedAxisRemover}
 
   // Export the Prime axis marker and the type classes that manipulate it
   export dimwit.prime.{Prime, PrimeRemover, PrimeRest, PrimeConcat}

core/src/main/scala/dimwit/random/Random.scala

@@ -73,7 +73,7 @@ object Random:
     *    Array(3, 4),
     *    Array(5, 6)) )
     *    val perm = Random.permutation(Axis[Row] -> 3)(key)
-    *    val shuffled = data.take(Axis[Row])(perm)
+    *    val shuffled = data(Axis[Row].at(perm))
     * }}}
     *
     * @param dim

core/src/main/scala/dimwit/tensor/Axis.scala

@@ -5,6 +5,8 @@ import dimwit.|*|
 
 import ShapeTypeHelpers.AxisIndex
 
+import scala.annotation.targetName
+
 /** Instances of this class represent an axis in a tensor with a specific label `L`.
   * Axis objects are used whenever an axis needs to be selected at the value level,
   * such as when indexing into a tensor or defining the shape of a tensor.
@@ -16,6 +18,10 @@ final class Axis[L: Label]:
   def at(range: Range): AxisAtRange[L] = AxisAtRange(this, range)
   def at(indices: Seq[Int]): AxisAtIndices[L] = AxisAtIndices(this, indices)
   def at(index: Tensor0[Int32]): AxisAtTensorIndex[L] = AxisAtTensorIndex(this, index)
+  @targetName("atTensorIndices")
+  def at[L2](indices: Tensor1[L2, Int32]): AxisAtTensorIndices[L, L2] = AxisAtTensorIndices(this, indices)
+  def at(start: Int, windowSize: Int): AxisAtWindow[L] = AxisAtWindow(this, start, windowSize)
+  def at(start: Tensor0[Int32], windowSize: Int): AxisAtWindow[L] = AxisAtWindow(this, start, windowSize)
   def at[I <: NonEmptyTuple](indices: I): AxisAtTupleIndices[L, I] = AxisAtTupleIndices(this, indices)
   def as[U](newAxis: Axis[U]): (Axis[L], Axis[U]) = (this, newAxis)
 
@@ -54,5 +60,15 @@ case class AxisAtIndices[L](axis: Axis[L], indices: Seq[Int]) extends AxisSelect
 /* Represent an axis selection by a tensor containing indices. This allows for dynamic indexing based on the contents of the tensor. */
 case class AxisAtTensorIndex[L](axis: Axis[L], index: Tensor0[Int32]) extends AxisSelector[L]
 
+/* Represent an axis selection by a vector of indices. The selected axis is replaced by the axis of the vector. */
+case class AxisAtTensorIndices[L, L2](axis: Axis[L], indices: Tensor1[L2, Int32]) extends AxisSelector[L]
+
+/** Represent an axis selection by a window of `windowSize` consecutive indices from `start`.
+  * The start may be a tensor, also a traced one (e.g. a loop index under `jit`); the size is static.
+  * A window with an `Int` start must fit the axis. With a tensor start, which cannot be checked,
+  * the start is clamped so that the window fits, as in `jax.lax.dynamic_slice`.
+  */
+case class AxisAtWindow[L](axis: Axis[L], start: Int | Tensor0[Int32], windowSize: Int) extends AxisSelector[L]
+
 /* Represent an axis selection by a tuple containing indices. */
 case class AxisAtTupleIndices[L, I <: NonEmptyTuple](axis: Axis[L], indices: I) extends AxisSelector[L]

core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala

@@ -164,22 +164,6 @@ object ShapeTypeHelpers:
       type RemainingAxes = R
       def indices: List[Int] = idx.indices
 
-  /** Removes [[RemovedAxis]] from a tensor shape while computing runtime indices
-    *  for [[IndexAxes]].
-    */
-  trait AxesConditionalRemover[TensorShape <: Tuple, RemovedAxis <: Tuple, IndexAxes <: Tuple] extends AxesInTensor[TensorShape, IndexAxes]:
-    type RemainingAxes <: Tuple
-
-  object AxesConditionalRemover:
-    type Aux[T <: Tuple, RA <: Tuple, IA <: Tuple, R <: Tuple] = AxesConditionalRemover[T, RA, IA] { type RemainingAxes = R }
-
-    given bridge[T <: Tuple, RemovedAxis <: Tuple, IndexAxes <: Tuple, R <: Tuple](using
-        idx: AxisIndices[T, IndexAxes],
-        ev: RemoverAll.Aux[T, RemovedAxis, R]
-    ): AxesConditionalRemover.Aux[T, RemovedAxis, IndexAxes, R] = new AxesConditionalRemover[T, RemovedAxis, IndexAxes]:
-      type RemainingAxes = R
-      def indices = idx.indices
-
   /** Removes a shared axis from multiple tensor shapes while computing runtime indices
     * for the remaining axes.
     */