Make updated chainable to write along several axes

AGENTS.md

@@ -667,7 +667,10 @@ val concatenated = concatenate(t1, t2, Axis[A])
 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.
+A selector selects along exactly one axis; to select along several axes, chain the selectors,
+for reading `t(s1)(s2)` and for writing `t.focus(s1).focus(s2).updated(value)`.
+`t.focus(selector)` focuses on the selection: `.updated(value)` writes it, `.updatedWith(f)` applies `f` to it
+and `.get` reads it; `t.updated(selector)(value)` is short for `t.focus(selector).updated(value)`.
 
 ```scala
 // Shape: A=4, B=3, the value at (a, b) is 10a + b
@@ -697,6 +700,14 @@ val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 1
 // 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))
+
+// Several axes: chain focus and write the value, which has the shape of grid(s1)(s2)
+val gridEntryWritten = grid.focus(Axis[A].at(1)).focus(Axis[B].at(2)).updated(Tensor0(-1.0f)) // entry (1, 2) is -1
+val gridBlockWritten = grid.focus(Axis[A].at(1 until 3)).focus(Axis[B].at(order)).updated(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 3)).fill(0.0f))
+
+// updatedWith applies a function to the selection, get reads it
+val gridEntryRead: Tensor0[Float32] = grid.focus(Axis[A].at(1)).focus(Axis[B].at(2)).get // 12
+val gridDoubledRow = grid.focus(Axis[A].at(1)).updatedWith(_ *! Tensor0(2.0f)) // row 1 is [20, 22, 24]
 ```
 
 The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one,
@@ -750,10 +761,10 @@ val gridWrongValue = grid.updated(Axis[A].at(0))(grid)
 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:737)
+// 	at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:762)
 // 	at dimwit.tensor.tensorops.StructuralExtensions$.select(StructuralOps.scala:600)
 // 	at dimwit.package$.select(package.scala:80)
-// 	at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:723)
+// 	at repl.MdocSession$MdocApp17.$init$$$anonfun$4(AGENTS.md:739)
 ```
 
 ---

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

@@ -90,6 +90,7 @@ package object dimwit:
   export dimwit.tensor.tensorops.StructuralOps.{where, where_!, triu, tril, stack, concatenate}
   export dimwit.tensor.tensorops.FunctionalOps.zipvmap
   export dimwit.tensor.tensorops.ControlFlowOps.{scan, foriLoop, whileLoop, cond}
+  export dimwit.tensor.tensorops.TensorFocus
 
   // Export convolution options
   export dimwit.tensor.{Padding, Stride1, Stride2, Stride3}

core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala

@@ -711,6 +711,31 @@ private[dimwit] object StructuralExtensions:
       requireSelectionShape(value, tensor.windowExtents(selector, axisIndex.index))
       Tensor(Jax.lax.dynamic_update_slice(tensor.jaxValue, value.jaxValue, tensor.windowStarts(selector, axisIndex.index)))
 
+    // `tensor.focus(selector)` is a [[TensorFocus]] on the selection: `.updated(value)` is `tensor.updated(selector)(value)`,
+    // `.get` is `tensor.select(selector)`, and `.focus` with a further selector selects along another axis of the selection,
+    // so `tensor.focus(s1).focus(s2).updated(value)` writes what `tensor(s1)(s2)` reads, e.g. a single entry.
+
+    def focus[L](selector: AxisAtIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorFocus[T, ev.RemainingAxes, V] =
+      TensorFocus(selected => tensor.select(selector)(using ev, selected), tensor.updated(selector)(using ev, labels)(_))
+
+    def focus[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorFocus[T, ev.RemainingAxes, V] =
+      TensorFocus(selected => tensor.select(selector)(using ev, selected), tensor.updated(selector)(using ev, labels)(_))
+
+    def focus[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorFocus[T, T, V] =
+      TensorFocus(_ => tensor.select(selector), tensor.updated(selector)(_))
+
+    def focus[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorFocus[T, T, V] =
+      TensorFocus(_ => tensor.select(selector), tensor.updated(selector)(_))
+
+    def focus[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorFocus[T, T, V] =
+      TensorFocus(_ => tensor.select(selector), tensor.updated(selector)(_))
+
+    def focus[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[T, L, L2], labels: Labels[T]): TensorFocus[T, ev.NewShape, V] =
+      TensorFocus(selected => tensor.select(selector)(using ev, selected), tensor.updated(selector)(using ev, labels)(_))
+
+    def focus[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorFocus[T, T, V] =
+      TensorFocus(_ => tensor.select(selector), tensor.updated(selector)(_))
+
     /** The Python index `tensor[:, ..., index, ..., :]`, with `index` at the axis `dimIndex`. */
     private def indexAlong(dimIndex: Int, index: py.Any): Jax.PyDynamic =
       val all = py.Dynamic.global.slice(py.None)

core/src/main/scala/dimwit/tensor/tensorops/TensorFocus.scala

@@ -0,0 +1,70 @@
+package dimwit.tensor.tensorops
+
+import dimwit.tensor.AxisAtIndex
+import dimwit.tensor.AxisAtIndices
+import dimwit.tensor.AxisAtRange
+import dimwit.tensor.AxisAtTensorIndex
+import dimwit.tensor.AxisAtTensorIndices
+import dimwit.tensor.AxisAtTupleIndices
+import dimwit.tensor.AxisAtWindow
+import dimwit.tensor.Labels
+import dimwit.tensor.ShapeTypeHelpers.AxisIndex
+import dimwit.tensor.ShapeTypeHelpers.AxisRemover
+import dimwit.tensor.ShapeTypeHelpers.AxisReplacer
+import dimwit.tensor.Tensor
+
+import StructuralExtensions.*
+
+/** A focus on a selection of a tensor, returned by `tensor.focus(selector)`, to read or write it.
+  *
+  * `.get` reads the selection, `.updated(value)` returns a copy of the tensor with `value` at the selection,
+  * and `.updatedWith(f)` returns a copy with `f` applied to the selection. `.focus` with a further selector selects
+  * along an axis of the selection, as `select` does, so that `tensor.focus(s1).focus(s2).updated(value)` writes
+  * `value` where `tensor(s1)(s2)` reads, e.g. a single entry. Each step reads its selection and writes it back.
+  *
+  * It behaves like a lens as long as the selected positions are distinct and in bounds. For repeated indices,
+  * it is unspecified which of their values is written; a traced index out of bounds is read clamped and not written.
+  *
+  * @tparam T the shape of the tensor
+  * @tparam S the shape of the selection
+  */
+final class TensorFocus[T <: Tuple, S <: Tuple, V] private[tensorops] (
+    private[tensorops] val selection: Labels[S] => Tensor[S, V],
+    write: Tensor[S, V] => Tensor[T, V]
+):
+
+  /** The selection, the same as the chained `select` calls. */
+  def get(using labels: Labels[S]): Tensor[S, V] = selection(labels)
+
+  /** The copy of the tensor with `value`, which has the shape of the selection, at the selection. */
+  def updated(value: Tensor[S, V]): Tensor[T, V] = write(value)
+
+  /** The copy of the tensor with `f` applied to the selection. `f` must keep the shape of the selection. */
+  def updatedWith(f: Tensor[S, V] => Tensor[S, V])(using labels: Labels[S]): Tensor[T, V] = write(f(selection(labels)))
+
+  // Focusing along a further axis writes back what it changes in the selection.
+
+  def focus[L](selector: AxisAtIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorFocus[T, ev.RemainingAxes, V] =
+    andThen(get.focus(selector)(using ev, labels))
+
+  def focus[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorFocus[T, ev.RemainingAxes, V] =
+    andThen(get.focus(selector)(using ev, labels))
+
+  def focus[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorFocus[T, S, V] =
+    andThen(get.focus(selector))
+
+  def focus[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorFocus[T, S, V] =
+    andThen(get.focus(selector))
+
+  def focus[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorFocus[T, S, V] =
+    andThen(get.focus(selector))
+
+  def focus[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[S, L, L2], labels: Labels[S]): TensorFocus[T, ev.NewShape, V] =
+    andThen(get.focus(selector)(using ev, labels))
+
+  def focus[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorFocus[T, S, V] =
+    andThen(get.focus(selector))
+
+  /** The focus on a selection of the selection: its value is written into the selection, and that into the tensor. */
+  private def andThen[S2 <: Tuple](inner: TensorFocus[S, S2, V]): TensorFocus[T, S2, V] =
+    TensorFocus(inner.selection, value => write(inner.updated(value)))

core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala

@@ -1,6 +1,7 @@
 package dimwit.tensor
 
 import dimwit.*
+import dimwit.autodiff.Autodiff
 import scala.compiletime.testing.typeCheckErrors
 
 class TensorOpsStructureSuite extends DimwitTest:
@@ -509,6 +510,76 @@ class TensorOpsStructureSuite extends DimwitTest:
         t.roll(Axis[B], shift = 1) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 1.0f, 2.0f), Array(6.0f, 4.0f, 5.0f))))
         t.roll(Axis[A], shift = 1) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(4.0f, 5.0f, 6.0f), Array(1.0f, 2.0f, 3.0f))))
 
+  describe("focus"):
+
+    // Shape: A=3, B=4, the value at (a, b) is 10a + b
+    val ab = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(3, 4)((a, b) => 10f * a + b))
+
+    it("updated along one axis is the updated of the tensor"):
+      val row = Tensor1(Axis[B]).fromArray(Array(0f, 0f, 0f, 0f))
+      ab.focus(Axis[A].at(1)).updated(row) should approxEqual(ab.updated(Axis[A].at(1))(row))
+
+    it("get is the chained select"):
+      ab.focus(Axis[A].at(1)).get should approxEqual(ab(Axis[A].at(1)))
+      ab.focus(Axis[A].at(1)).focus(Axis[B].at(2)).get.item shouldBe 12f
+
+    it("updatedWith applies a function to the selection"):
+      ab.focus(Axis[A].at(1)).updatedWith(_ *! 2f) should approxEqual(ab.updated(Axis[A].at(1))(ab(Axis[A].at(1)) *! 2f))
+      ab.focus(Axis[A].at(1)).focus(Axis[B].at(2)).updatedWith(_ + Tensor0(1f))(Axis[A].at(1))(Axis[B].at(2)).item shouldBe 13f
+
+    it("writes a single entry"):
+      val written = ab.focus(Axis[A].at(1)).focus(Axis[B].at(2)).updated(Tensor0(-1f))
+      written should approxEqual(ab.updated(Axis[A].at(1))(ab(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1f))))
+      written(Axis[A].at(1))(Axis[B].at(2)).item shouldBe -1f
+      ab.focus(Axis[B].at(2)).focus(Axis[A].at(1)).updated(-1f) should approxEqual(written)
+
+    it("writes where the chained select reads, for every selector"):
+      val idx = Tensor1(Axis[C]).fromArray(Array(3, 0))
+      ab.focus(Axis[A].at(0 to 2 by 2)).focus(Axis[B].at(idx)).updated(Tensor(Shape(Axis[A] -> 2, Axis[C] -> 2)).fill(-1f)) should approxEqual(
+        Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(-1f, 1f, 2f, -1f), Array(10f, 11f, 12f, 13f), Array(-1f, 21f, 22f, -1f)))
+      )
+      val written = ab
+        .focus(Axis[A].at(Seq(2, 1)))
+        .focus(Axis[A].at((1, 0)))
+        .focus(Axis[B].at(Tensor0(1), 2))
+        .focus(Axis[A].at(Tensor0(0)))
+        .focus(Axis[B].at(1 until 2))
+        .updated(Tensor1(Axis[B]).fromArray(Array(-1f)))
+      // rows 2, 1 -> rows 1, 2 -> row 1; columns 1, 2 -> column 2
+      written(Axis[A].at(1))(Axis[B].at(2)).item shouldBe -1f
+      written.sum.item shouldBe ab.sum.item - 12f - 1f
+
+    it("writes three levels deep"):
+      val abc = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3, Axis[C] -> 4)).fill(0f)
+      val written = abc.focus(Axis[C].at(3)).focus(Axis[A].at(1)).focus(Axis[B].at(0)).updated(Tensor0(1f))
+      written(Axis[A].at(1))(Axis[B].at(0))(Axis[C].at(3)).item shouldBe 1f
+      written.sum.item shouldBe 1f
+
+    it("writes entries at traced indices under jit"):
+      val filled = jit((t: Tensor2[A, B, Float32]) =>
+        foriLoop(0, 3)(t): (i, acc) =>
+          foriLoop(0, 4)(acc): (j, acc) =>
+            acc.focus(Axis[A].at(i)).focus(Axis[B].at(j)).updated((i * Tensor0(10) + j).asType(VType[Float32]))
+      )(Tensor(ab.shape).fill(0f))
+      filled should approxEqual(ab)
+
+    it("is differentiable"):
+      val grad = Autodiff.grad((v: Tensor0[Float32]) => (ab.focus(Axis[A].at(1)).focus(Axis[B].at(2)).updated(v) * ab).sum)
+      grad(Tensor0(0f)).value.item shouldBe 12f
+
+    it("updated writes directly and does not chain"):
+      "ab.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(0f))" shouldNot compile
+      "ab.focus(Axis[A].at(1))(Tensor1(Axis[B]).fromArray(Array(0f, 0f, 0f, 0f)))" shouldNot compile
+
+    it("rejects an axis that the selection does not have"):
+      "ab.focus(Axis[A].at(1)).focus(Axis[A].at(0)).updated(Tensor0(0f))" shouldNot compile
+      "ab.focus(Axis[A].at(1)).focus(Axis[C].at(0)).updated(Tensor0(0f))" shouldNot compile
+      "ab.focus(Axis[A].at(1)).focus(Axis[B].at(0)).updated(Tensor1(Axis[B]).fromArray(Array(0f)))" shouldNot compile
+
+    it("rejects a value whose shape differs from the selection"):
+      the[IllegalArgumentException] thrownBy ab.focus(Axis[B].at(0)).focus(Axis[A].at(0 to 1)).updated(Tensor1(Axis[A]).fromArray(Array(0f, 0f, 0f))) should have message
+        "requirement failed: Cannot set a value of shape Shape(A -> 3) into a selection of shape (A -> 2)"
+
   describe("window (AxisAtWindow)"):
 
     // Shape: A=2, B=5, C=3, the value at (a, b, c) is 100a + 10b + c

docs/quickstart.md

@@ -334,6 +334,12 @@ as `updated` does for Scala's immutable collections:
   val updated1 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1))(Tensor(Shape(Axis[B] -> 2)).fill(0.0f))
 ```
 
+To write along several axes, e.g. a single entry, focus on the selection one axis at a time and write the value:
+
+```scala
+  val updated2 : Tensor2[A, B, Float32] = concatenated.focus(Axis[A].at(1)).focus(Axis[B].at(0)).updated(Tensor0(0.0f))
+```
+
 #### Squeezing, Expanding and transposing axes
 
 Given a tensor, for which one axis has extent 1, we can remove that axis using the `squeeze`

mdocs/AGENTS.md

@@ -487,7 +487,10 @@ val concatenated = concatenate(t1, t2, Axis[A])
 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.
+A selector selects along exactly one axis; to select along several axes, chain the selectors,
+for reading `t(s1)(s2)` and for writing `t.focus(s1).focus(s2).updated(value)`.
+`t.focus(selector)` focuses on the selection: `.updated(value)` writes it, `.updatedWith(f)` applies `f` to it
+and `.get` reads it; `t.updated(selector)(value)` is short for `t.focus(selector).updated(value)`.
 
 ```scala mdoc:silent
 // Shape: A=4, B=3, the value at (a, b) is 10a + b
@@ -517,6 +520,14 @@ val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 1
 // 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))
+
+// Several axes: chain focus and write the value, which has the shape of grid(s1)(s2)
+val gridEntryWritten = grid.focus(Axis[A].at(1)).focus(Axis[B].at(2)).updated(Tensor0(-1.0f)) // entry (1, 2) is -1
+val gridBlockWritten = grid.focus(Axis[A].at(1 until 3)).focus(Axis[B].at(order)).updated(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 3)).fill(0.0f))
+
+// updatedWith applies a function to the selection, get reads it
+val gridEntryRead: Tensor0[Float32] = grid.focus(Axis[A].at(1)).focus(Axis[B].at(2)).get // 12
+val gridDoubledRow = grid.focus(Axis[A].at(1)).updatedWith(_ *! Tensor0(2.0f)) // row 1 is [20, 22, 24]
 ```
 
 The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one,

mdocs/docs/quickstart.md

@@ -322,6 +322,12 @@ as `updated` does for Scala's immutable collections:
   val updated1 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1))(Tensor(Shape(Axis[B] -> 2)).fill(0.0f))
 ```
 
+To write along several axes, e.g. a single entry, focus on the selection one axis at a time and write the value:
+
+```scala mdoc:silent
+  val updated2 : Tensor2[A, B, Float32] = concatenated.focus(Axis[A].at(1)).focus(Axis[B].at(0)).updated(Tensor0(0.0f))
+```
+
 #### Squeezing, Expanding and transposing axes
 
 Given a tensor, for which one axis has extent 1, we can remove that axis using the `squeeze`