diff --git a/AGENTS.md b/AGENTS.md index 118e2c2..9fcb294 100644 --- a/AGENTS.md +++ b/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) ``` --- diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 75fb927..a3a4d08 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/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} diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index abad3d6..60aaced 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/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) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/TensorFocus.scala b/core/src/main/scala/dimwit/tensor/tensorops/TensorFocus.scala new file mode 100644 index 0000000..e42bc97 --- /dev/null +++ b/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))) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index a084c72..86d6fe1 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/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 diff --git a/docs/quickstart.md b/docs/quickstart.md index 96cacf8..1048d7d 100644 --- a/docs/quickstart.md +++ b/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` diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 91e190f..919b1b8 100644 --- a/mdocs/AGENTS.md +++ b/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, diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md index d6854bd..73f2d13 100644 --- a/mdocs/docs/quickstart.md +++ b/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`