From 8bf693091ff11624510898bccc505b95792818fd Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 24 Oct 2026 08:48:44 +0200 Subject: [PATCH 1/2] Make updated chainable to write along several axes So far, updated writes along one axis and directly returns the updated tensor, so writing along several axes at once, e.g. a single entry, needs a manual read, modify and write back. updated now returns a TensorUpdate. Applied to a value, it returns the updated tensor as before, so existing code is unchanged. Its updated method takes a further selector, which allows chaining updates like chained select: t.updated(Axis[A].at(1)).updated(Axis[B].at(2))(3f) updates the entry at (1, 2). Each step reads the selection and writes it back, reusing select and updated. In the future, TensorUpdate could support other operations besides setting the values, such as adding them or taking minimums, like JAX's x.at[...]. --- AGENTS.md | 15 ++- core/src/main/scala/dimwit/package.scala | 1 + .../tensor/tensorops/StructuralOps.scala | 125 ++++++++++-------- .../tensor/tensorops/TensorUpdate.scala | 60 +++++++++ .../tensor/TensorOpsStructureSuite.scala | 58 ++++++++ docs/quickstart.md | 6 + mdocs/AGENTS.md | 7 +- mdocs/docs/quickstart.md | 6 + 8 files changed, 215 insertions(+), 63 deletions(-) create mode 100644 core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala diff --git a/AGENTS.md b/AGENTS.md index a5cb9ad..4543515 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -667,7 +667,8 @@ 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.updated(s1).updated(s2)(value)`. ```scala // Shape: A=4, B=3, the value at (a, b) is 10a + b @@ -697,6 +698,10 @@ 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 updated; the value has the shape of grid(s1)(s2) +val gridEntryWritten = grid.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1.0f)) // entry (1, 2) is -1 +val gridBlockWritten = grid.updated(Axis[A].at(1 until 3)).updated(Axis[B].at(order))(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 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, @@ -735,9 +740,9 @@ val gridWrongValue = grid.updated(Axis[A].at(0))(grid) // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, // dimwit.tensor.DType.Float32] // ) -// Required: dimwit.tensor.Tensor[?1.RemainingAxes, dimwit.tensor.DType.Float32] +// Required: dimwit.tensor.Tensor[?2.RemainingAxes, dimwit.tensor.DType.Float32] // -// where: ?1 is an unknown value of type dimwit.tensor.ShapeTypeHelpers.AxisRemover.Aux[ +// where: ?2 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] // @@ -750,10 +755,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:742) +// at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:753) // 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) +// at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:731) ``` --- diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 75fb927..c26fe50 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.TensorUpdate // 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 eaa96f2..a5f878c 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -628,93 +628,104 @@ private[dimwit] object StructuralExtensions: def apply[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): Tensor[T, V] = tensor.select(selector) - /** A copy of the tensor with new values at an index along an axis. + // `tensor.updated(selector)(value)` returns a copy of the tensor with `value` at the selection, as `updated` does for + // Scala's immutable collections. `tensor.updated(selector)` is a [[TensorUpdate]]: its `updated` with a further selector + // selects along another axis of the selection, so `tensor.updated(s1).updated(s2)(value)` writes what `tensor(s1)(s2)` reads. + + /** Update the tensor at an index along an axis. * * @param selector of the form `Axis[L].at(index)` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtIndex[L])(using - ev: AxisRemover[T, L], - labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index))).set(value.jaxValue)) + def updated[L](selector: AxisAtIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorUpdate[T, ev.RemainingAxes, V] = + TensorUpdate( + selected => tensor.select(selector)(using ev, selected), + value => + requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index))).set(value.jaxValue)) + ) - /** A copy of the tensor with new values at an index given by a tensor along an axis. The index may be traced. + /** Update the tensor at an index given by a tensor along an axis. The index may be traced. * An out-of-bounds index leaves the tensor unchanged. * * @param selector of the form `Axis[L].at(index: Tensor0[Int32])` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtTensorIndex[L])(using - ev: AxisRemover[T, L], - labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue)).set(value.jaxValue)) + def updated[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorUpdate[T, ev.RemainingAxes, V] = + TensorUpdate( + selected => tensor.select(selector)(using ev, selected), + value => + requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue)).set(value.jaxValue)) + ) - /** A copy of the tensor with new values at a range along an axis. + /** Update the tensor at a range along an axis. * * @param selector of the form `Axis[L].at(range)` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtRange[L])(using - axisIndex: AxisIndex[T, L], - labels: Labels[T] - )(value: Tensor[T, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.range.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index))).set(value.jaxValue)) + def updated[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = + TensorUpdate( + _ => tensor.select(selector), + value => + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.range.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index))).set(value.jaxValue)) + ) - /** A copy of the tensor with new values at a sequence of indices along an axis. + /** Update the tensor at a sequence of indices along an axis. * * @param selector of the form `Axis[L].at(Seq(i, j, ...))` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtIndices[L])(using - axisIndex: AxisIndex[T, L], - labels: Labels[T] - )(value: Tensor[T, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.indices.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices))).set(value.jaxValue)) + def updated[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = + TensorUpdate( + _ => tensor.select(selector), + value => + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.indices.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices))).set(value.jaxValue)) + ) - /** A copy of the tensor with new values at a tuple of indices along an axis. + /** Update the tensor at a tuple of indices along an axis. * * @param selector of the form `Axis[L].at((i, j, ...))` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using - axisIndex: AxisIndex[T, L], - labels: Labels[T] - )(value: Tensor[T, V]): Tensor[T, V] = - val indices = selector.indices.toList.asInstanceOf[List[Int]] - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, indices.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(indices))).set(value.jaxValue)) + def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = + TensorUpdate( + _ => tensor.select(selector), + value => + val indices = selector.indices.toList.asInstanceOf[List[Int]] + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, indices.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(indices))).set(value.jaxValue)) + ) - /** A copy of the tensor with new values at the indices of a vector along an axis. Out-of-bounds indices are skipped; + /** Update the tensor at the indices of a vector along an axis. Out-of-bounds indices are skipped; * for repeated indices, it is unspecified which of their values is written. * * @param selector of the form `Axis[L].at(indices: Tensor1[L2, Int32])` - * @param value the new values, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` */ - def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using - ev: AxisReplacer[T, L, L2], - labels: Labels[T] - )(value: Tensor[ev.NewShape, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.shape.dimensions.updated(ev.index, selector.indices.shape.dimensions.head)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.indices.jaxValue)).set(value.jaxValue)) + def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[T, L, L2], labels: Labels[T]): TensorUpdate[T, ev.NewShape, V] = + TensorUpdate( + selected => tensor.select(selector)(using ev, selected), + value => + requireSelectionShape(value, tensor.shape.dimensions.updated(ev.index, selector.indices.shape.dimensions.head)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.indices.jaxValue)).set(value.jaxValue)) + ) - /** A copy of the tensor with new values in a window of static size along an axis, like `jax.lax.dynamic_update_slice`. + /** Update the tensor in a window of static size along an axis, like `jax.lax.dynamic_update_slice`. * Its start may be a tensor, also a traced one, e.g. a loop index under `jit`. * * @param selector of the form `Axis[L].at(start, windowSize)` - * @param value the new window, with the shape of `tensor.select(selector)` + * @return the update, to apply to the new window, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtWindow[L])(using - axisIndex: AxisIndex[T, L], - labels: Labels[T] - )(value: Tensor[T, V]): Tensor[T, V] = - requireSelectionShape(value, tensor.windowExtents(selector, axisIndex.index)) - Tensor(Jax.lax.dynamic_update_slice(tensor.jaxValue, value.jaxValue, tensor.windowStarts(selector, axisIndex.index))) + def updated[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = + TensorUpdate( + _ => tensor.select(selector), + value => + requireSelectionShape(value, tensor.windowExtents(selector, axisIndex.index)) + Tensor(Jax.lax.dynamic_update_slice(tensor.jaxValue, value.jaxValue, tensor.windowStarts(selector, axisIndex.index))) + ) /** The Python index `tensor[:, ..., index, ..., :]`, with `index` at the axis `dimIndex`. */ private def indexAlong(dimIndex: Int, index: py.Any): Jax.PyDynamic = diff --git a/core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala b/core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala new file mode 100644 index 0000000..42630ab --- /dev/null +++ b/core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala @@ -0,0 +1,60 @@ +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.* + +/** The update of a selection of a tensor, returned by `tensor.updated(selector)`. + * + * Applied to a value with the shape of the selection, it returns the updated copy of the tensor. + * `updated` with a further selector selects along an axis of the selection, as `select` does, so that + * `tensor.updated(s1).updated(s2)(value)` writes `value` where `tensor(s1)(s2)` reads, e.g. a single entry. + * + * @tparam T the shape of the tensor + * @tparam S the shape of the selection + */ +final class TensorUpdate[T <: Tuple, S <: Tuple, V] private[tensorops] ( + private[tensorops] val selection: Labels[S] => Tensor[S, V], + write: Tensor[S, V] => Tensor[T, V] +): + + /** The copy of the tensor with `value` at the selection. */ + def apply(value: Tensor[S, V]): Tensor[T, V] = write(value) + + // Selecting along a further axis writes back what it changes in the selection. + + def updated[L](selector: AxisAtIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorUpdate[T, ev.RemainingAxes, V] = + andThen(selection(labels).updated(selector)(using ev, labels)) + + def updated[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorUpdate[T, ev.RemainingAxes, V] = + andThen(selection(labels).updated(selector)(using ev, labels)) + + def updated[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = + andThen(selection(labels).updated(selector)) + + def updated[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = + andThen(selection(labels).updated(selector)) + + def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = + andThen(selection(labels).updated(selector)) + + def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[S, L, L2], labels: Labels[S]): TensorUpdate[T, ev.NewShape, V] = + andThen(selection(labels).updated(selector)(using ev, labels)) + + def updated[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = + andThen(selection(labels).updated(selector)) + + /** The update of a selection of the selection: its value is written into the selection, and that into the tensor. */ + private def andThen[S2 <: Tuple](inner: TensorUpdate[S, S2, V]): TensorUpdate[T, S2, V] = + TensorUpdate(inner.selection, value => write(inner(value))) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 190bd43..2896bcd 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: @@ -627,6 +628,63 @@ 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("updated along several axes"): + + // 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("writes a single entry"): + val written = ab.updated(Axis[A].at(1)).updated(Axis[B].at(2))(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.updated(Axis[B].at(2)).updated(Axis[A].at(1))(-1f) should approxEqual(written) + + it("writes where the chained select reads, for every selector"): + val idx = Tensor1(Axis[C]).fromArray(Array(3, 0)) + val selectors = ab.updated(Axis[A].at(0 to 2 by 2)).updated(Axis[B].at(idx)) + selectors(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 + .updated(Axis[A].at(Seq(2, 1))) + .updated(Axis[A].at((1, 0))) + .updated(Axis[B].at(Tensor0(1), 2)) + .updated(Axis[A].at(Tensor0(0))) + .updated(Axis[B].at(1 until 2))(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.updated(Axis[C].at(3)).updated(Axis[A].at(1)).updated(Axis[B].at(0))(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.updated(Axis[A].at(i)).updated(Axis[B].at(j))((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.updated(Axis[A].at(1)).updated(Axis[B].at(2))(v) * ab).sum) + grad(Tensor0(0f)).value.item shouldBe 12f + + it("takes further selectors with updated, not apply"): + "ab.updated(Axis[A].at(1))(Axis[B].at(2))(Tensor0(0f))" shouldNot compile + + it("rejects an axis that the selection does not have"): + "ab.updated(Axis[A].at(1)).updated(Axis[A].at(0))(Tensor0(0f))" shouldNot compile + "ab.updated(Axis[A].at(1)).updated(Axis[C].at(0))(Tensor0(0f))" shouldNot compile + "ab.updated(Axis[A].at(1)).updated(Axis[B].at(0))(Tensor1(Axis[B]).fromArray(Array(0f)))" shouldNot compile + + it("rejects a value whose shape differs from the selection"): + the[IllegalArgumentException] thrownBy ab.updated(Axis[B].at(0)).updated(Axis[A].at(0 to 1))(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..f5d150b 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, chain `updated` with one selector per axis: + +```scala + val updated2 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1)).updated(Axis[B].at(0))(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..10c0eb3 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -487,7 +487,8 @@ 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.updated(s1).updated(s2)(value)`. ```scala mdoc:silent // Shape: A=4, B=3, the value at (a, b) is 10a + b @@ -517,6 +518,10 @@ 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 updated; the value has the shape of grid(s1)(s2) +val gridEntryWritten = grid.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1.0f)) // entry (1, 2) is -1 +val gridBlockWritten = grid.updated(Axis[A].at(1 until 3)).updated(Axis[B].at(order))(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 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, diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md index d6854bd..fe1bfdd 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, chain `updated` with one selector per axis: + +```scala mdoc:silent + val updated2 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1)).updated(Axis[B].at(0))(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` From 4d02abed9b197b52c0843b8dcdab746796f5a1e1 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Mon, 26 Oct 2026 08:25:59 +0100 Subject: [PATCH 2/2] introduce focus mechanism Instead of having chained update, this commit introduces the method focus, which, like a lens, focuses on some parts of the tensor and eventually makes it possible to update, read, modify, ... the parts on which the focus was set. --- AGENTS.md | 22 ++- core/src/main/scala/dimwit/package.scala | 2 +- .../tensor/tensorops/StructuralOps.scala | 150 ++++++++++-------- .../dimwit/tensor/tensorops/TensorFocus.scala | 70 ++++++++ .../tensor/tensorops/TensorUpdate.scala | 60 ------- .../tensor/TensorOpsStructureSuite.scala | 51 +++--- docs/quickstart.md | 4 +- mdocs/AGENTS.md | 14 +- mdocs/docs/quickstart.md | 4 +- 9 files changed, 213 insertions(+), 164 deletions(-) create mode 100644 core/src/main/scala/dimwit/tensor/tensorops/TensorFocus.scala delete mode 100644 core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala diff --git a/AGENTS.md b/AGENTS.md index 4543515..1985ded 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -668,7 +668,9 @@ Every selection along an axis is an `Axis[L].at(...)` selector. `select` reads t 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 selectors, -for reading `t(s1)(s2)` and for writing `t.updated(s1).updated(s2)(value)`. +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 @@ -699,9 +701,13 @@ val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 1 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 updated; the value has the shape of grid(s1)(s2) -val gridEntryWritten = grid.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1.0f)) // entry (1, 2) is -1 -val gridBlockWritten = grid.updated(Axis[A].at(1 until 3)).updated(Axis[B].at(order))(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 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, @@ -740,9 +746,9 @@ val gridWrongValue = grid.updated(Axis[A].at(0))(grid) // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, // dimwit.tensor.DType.Float32] // ) -// Required: dimwit.tensor.Tensor[?2.RemainingAxes, dimwit.tensor.DType.Float32] +// Required: dimwit.tensor.Tensor[?1.RemainingAxes, dimwit.tensor.DType.Float32] // -// where: ?2 is an unknown value of type dimwit.tensor.ShapeTypeHelpers.AxisRemover.Aux[ +// 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] // @@ -755,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:753) +// at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:767) // 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:731) +// 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 c26fe50..a3a4d08 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -90,7 +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.TensorUpdate + 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 a5f878c..d120bb1 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -628,104 +628,118 @@ private[dimwit] object StructuralExtensions: def apply[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): Tensor[T, V] = tensor.select(selector) - // `tensor.updated(selector)(value)` returns a copy of the tensor with `value` at the selection, as `updated` does for - // Scala's immutable collections. `tensor.updated(selector)` is a [[TensorUpdate]]: its `updated` with a further selector - // selects along another axis of the selection, so `tensor.updated(s1).updated(s2)(value)` writes what `tensor(s1)(s2)` reads. - - /** Update the tensor at an index along an axis. + /** A copy of the tensor with new values at an index along an axis. * * @param selector of the form `Axis[L].at(index)` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorUpdate[T, ev.RemainingAxes, V] = - TensorUpdate( - selected => tensor.select(selector)(using ev, selected), - value => - requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index))).set(value.jaxValue)) - ) + def updated[L](selector: AxisAtIndex[L])(using + ev: AxisRemover[T, L], + labels: Labels[T] + )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = + requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index))).set(value.jaxValue)) - /** Update the tensor at an index given by a tensor along an axis. The index may be traced. + /** A copy of the tensor with new values at an index given by a tensor along an axis. The index may be traced. * An out-of-bounds index leaves the tensor unchanged. * * @param selector of the form `Axis[L].at(index: Tensor0[Int32])` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[T, L], labels: Labels[T]): TensorUpdate[T, ev.RemainingAxes, V] = - TensorUpdate( - selected => tensor.select(selector)(using ev, selected), - value => - requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue)).set(value.jaxValue)) - ) + def updated[L](selector: AxisAtTensorIndex[L])(using + ev: AxisRemover[T, L], + labels: Labels[T] + )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = + requireSelectionShape(value, tensor.shape.dimensions.patch(ev.index, Nil, 1)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue)).set(value.jaxValue)) - /** Update the tensor at a range along an axis. + /** A copy of the tensor with new values at a range along an axis. * * @param selector of the form `Axis[L].at(range)` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = - TensorUpdate( - _ => tensor.select(selector), - value => - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.range.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index))).set(value.jaxValue)) - ) + def updated[L](selector: AxisAtRange[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + )(value: Tensor[T, V]): Tensor[T, V] = + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.range.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index))).set(value.jaxValue)) - /** Update the tensor at a sequence of indices along an axis. + /** A copy of the tensor with new values at a sequence of indices along an axis. * * @param selector of the form `Axis[L].at(Seq(i, j, ...))` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = - TensorUpdate( - _ => tensor.select(selector), - value => - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.indices.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices))).set(value.jaxValue)) - ) + def updated[L](selector: AxisAtIndices[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + )(value: Tensor[T, V]): Tensor[T, V] = + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, selector.indices.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices))).set(value.jaxValue)) - /** Update the tensor at a tuple of indices along an axis. + /** A copy of the tensor with new values at a tuple of indices along an axis. * * @param selector of the form `Axis[L].at((i, j, ...))` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = - TensorUpdate( - _ => tensor.select(selector), - value => - val indices = selector.indices.toList.asInstanceOf[List[Int]] - requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, indices.length)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(indices))).set(value.jaxValue)) - ) + def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + )(value: Tensor[T, V]): Tensor[T, V] = + val indices = selector.indices.toList.asInstanceOf[List[Int]] + requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, indices.length)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(indices))).set(value.jaxValue)) - /** Update the tensor at the indices of a vector along an axis. Out-of-bounds indices are skipped; + /** A copy of the tensor with new values at the indices of a vector along an axis. Out-of-bounds indices are skipped; * for repeated indices, it is unspecified which of their values is written. * * @param selector of the form `Axis[L].at(indices: Tensor1[L2, Int32])` - * @return the update, to apply to the new values, with the shape of `tensor.select(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[T, L, L2], labels: Labels[T]): TensorUpdate[T, ev.NewShape, V] = - TensorUpdate( - selected => tensor.select(selector)(using ev, selected), - value => - requireSelectionShape(value, tensor.shape.dimensions.updated(ev.index, selector.indices.shape.dimensions.head)) - Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.indices.jaxValue)).set(value.jaxValue)) - ) + def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using + ev: AxisReplacer[T, L, L2], + labels: Labels[T] + )(value: Tensor[ev.NewShape, V]): Tensor[T, V] = + requireSelectionShape(value, tensor.shape.dimensions.updated(ev.index, selector.indices.shape.dimensions.head)) + Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(ev.index, selector.indices.jaxValue)).set(value.jaxValue)) - /** Update the tensor in a window of static size along an axis, like `jax.lax.dynamic_update_slice`. + /** A copy of the tensor with new values in a window of static size along an axis, like `jax.lax.dynamic_update_slice`. * Its start may be a tensor, also a traced one, e.g. a loop index under `jit`. * * @param selector of the form `Axis[L].at(start, windowSize)` - * @return the update, to apply to the new window, with the shape of `tensor.select(selector)` + * @param value the new window, with the shape of `tensor.select(selector)` */ - def updated[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): TensorUpdate[T, T, V] = - TensorUpdate( - _ => tensor.select(selector), - value => - requireSelectionShape(value, tensor.windowExtents(selector, axisIndex.index)) - Tensor(Jax.lax.dynamic_update_slice(tensor.jaxValue, value.jaxValue, tensor.windowStarts(selector, axisIndex.index))) - ) + def updated[L](selector: AxisAtWindow[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + )(value: Tensor[T, V]): Tensor[T, V] = + 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 = 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/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala b/core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala deleted file mode 100644 index 42630ab..0000000 --- a/core/src/main/scala/dimwit/tensor/tensorops/TensorUpdate.scala +++ /dev/null @@ -1,60 +0,0 @@ -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.* - -/** The update of a selection of a tensor, returned by `tensor.updated(selector)`. - * - * Applied to a value with the shape of the selection, it returns the updated copy of the tensor. - * `updated` with a further selector selects along an axis of the selection, as `select` does, so that - * `tensor.updated(s1).updated(s2)(value)` writes `value` where `tensor(s1)(s2)` reads, e.g. a single entry. - * - * @tparam T the shape of the tensor - * @tparam S the shape of the selection - */ -final class TensorUpdate[T <: Tuple, S <: Tuple, V] private[tensorops] ( - private[tensorops] val selection: Labels[S] => Tensor[S, V], - write: Tensor[S, V] => Tensor[T, V] -): - - /** The copy of the tensor with `value` at the selection. */ - def apply(value: Tensor[S, V]): Tensor[T, V] = write(value) - - // Selecting along a further axis writes back what it changes in the selection. - - def updated[L](selector: AxisAtIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorUpdate[T, ev.RemainingAxes, V] = - andThen(selection(labels).updated(selector)(using ev, labels)) - - def updated[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[S, L], labels: Labels[S]): TensorUpdate[T, ev.RemainingAxes, V] = - andThen(selection(labels).updated(selector)(using ev, labels)) - - def updated[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = - andThen(selection(labels).updated(selector)) - - def updated[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = - andThen(selection(labels).updated(selector)) - - def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = - andThen(selection(labels).updated(selector)) - - def updated[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[S, L, L2], labels: Labels[S]): TensorUpdate[T, ev.NewShape, V] = - andThen(selection(labels).updated(selector)(using ev, labels)) - - def updated[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[S, L], labels: Labels[S]): TensorUpdate[T, S, V] = - andThen(selection(labels).updated(selector)) - - /** The update of a selection of the selection: its value is written into the selection, and that into the tensor. */ - private def andThen[S2 <: Tuple](inner: TensorUpdate[S, S2, V]): TensorUpdate[T, S2, V] = - TensorUpdate(inner.selection, value => write(inner(value))) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 2896bcd..81f4cc5 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -628,36 +628,48 @@ 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("updated along several axes"): + 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.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1f)) + 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.updated(Axis[B].at(2)).updated(Axis[A].at(1))(-1f) should approxEqual(written) + 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)) - val selectors = ab.updated(Axis[A].at(0 to 2 by 2)).updated(Axis[B].at(idx)) - selectors(Tensor(Shape(Axis[A] -> 2, Axis[C] -> 2)).fill(-1f)) should approxEqual( + 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 - .updated(Axis[A].at(Seq(2, 1))) - .updated(Axis[A].at((1, 0))) - .updated(Axis[B].at(Tensor0(1), 2)) - .updated(Axis[A].at(Tensor0(0))) - .updated(Axis[B].at(1 until 2))(Tensor1(Axis[B]).fromArray(Array(-1f))) + .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.updated(Axis[C].at(3)).updated(Axis[A].at(1)).updated(Axis[B].at(0))(Tensor0(1f)) + 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 @@ -665,24 +677,25 @@ class TensorOpsStructureSuite extends DimwitTest: val filled = jit((t: Tensor2[A, B, Float32]) => foriLoop(0, 3)(t): (i, acc) => foriLoop(0, 4)(acc): (j, acc) => - acc.updated(Axis[A].at(i)).updated(Axis[B].at(j))((i * Tensor0(10) + j).asType(VType[Float32])) + 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.updated(Axis[A].at(1)).updated(Axis[B].at(2))(v) * ab).sum) + 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("takes further selectors with updated, not apply"): - "ab.updated(Axis[A].at(1))(Axis[B].at(2))(Tensor0(0f))" shouldNot compile + 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.updated(Axis[A].at(1)).updated(Axis[A].at(0))(Tensor0(0f))" shouldNot compile - "ab.updated(Axis[A].at(1)).updated(Axis[C].at(0))(Tensor0(0f))" shouldNot compile - "ab.updated(Axis[A].at(1)).updated(Axis[B].at(0))(Tensor1(Axis[B]).fromArray(Array(0f)))" shouldNot compile + "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.updated(Axis[B].at(0)).updated(Axis[A].at(0 to 1))(Tensor1(Axis[A]).fromArray(Array(0f, 0f, 0f))) should have message + 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)"): diff --git a/docs/quickstart.md b/docs/quickstart.md index f5d150b..1048d7d 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -334,10 +334,10 @@ 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, chain `updated` with one selector per axis: +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.updated(Axis[A].at(1)).updated(Axis[B].at(0))(Tensor0(0.0f)) + 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 diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 10c0eb3..919b1b8 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -488,7 +488,9 @@ Every selection along an axis is an `Axis[L].at(...)` selector. `select` reads t 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 selectors, -for reading `t(s1)(s2)` and for writing `t.updated(s1).updated(s2)(value)`. +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 @@ -519,9 +521,13 @@ val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 1 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 updated; the value has the shape of grid(s1)(s2) -val gridEntryWritten = grid.updated(Axis[A].at(1)).updated(Axis[B].at(2))(Tensor0(-1.0f)) // entry (1, 2) is -1 -val gridBlockWritten = grid.updated(Axis[A].at(1 until 3)).updated(Axis[B].at(order))(Tensor(Shape2(Axis[A] -> 2, Axis[C] -> 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 fe1bfdd..73f2d13 100644 --- a/mdocs/docs/quickstart.md +++ b/mdocs/docs/quickstart.md @@ -322,10 +322,10 @@ 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, chain `updated` with one selector per axis: +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.updated(Axis[A].at(1)).updated(Axis[B].at(0))(Tensor0(0.0f)) + 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