From 8dab3754d7f49d093f64d909b0ed954b8ea6b8f9 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Wed, 7 Oct 2026 17:15:49 +0200 Subject: [PATCH 1/3] Unify slicing on single-axis Axis[L].at selectors: remove the tuple form and take (now slice(Axis[L].at(indices))), and add the window selector Axis[L].at(start, windowSize), replacing Tensor1.dynamicSlice. set supports every selector and requires the value to have the shape of the selection. --- core/src/main/scala/dimwit/package.scala | 8 +- .../src/main/scala/dimwit/random/Random.scala | 2 +- core/src/main/scala/dimwit/tensor/Axis.scala | 16 + .../dimwit/tensor/ShapeTypeHelpers.scala | 16 - .../tensor/tensorops/StructuralOps.scala | 369 +++++++++--------- .../tensor/tensorops/Tensor1Extensions.scala | 17 - .../scala/dimwit/random/RandomSuite.scala | 2 +- .../dimwit/tensor/TensorCreationSuite.scala | 4 +- .../tensor/TensorOpsBroadcastSuite.scala | 24 +- .../tensor/TensorOpsStructureSuite.scala | 130 +++++- .../src/main/scala/dimwit/basic/KMeans.scala | 2 +- .../dimwit/basic/LogisticRegression.scala | 8 +- 12 files changed, 347 insertions(+), 251 deletions(-) diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 87ac333f..a5d33c0e 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -4,6 +4,8 @@ import dimwit.tensor.AxisAtIndex import dimwit.tensor.AxisAtIndices import dimwit.tensor.AxisAtRange import dimwit.tensor.AxisAtTensorIndex +import dimwit.tensor.AxisAtTensorIndices +import dimwit.tensor.AxisAtWindow import dimwit.tensor.AxisExtent import dimwit.tensor.AxisSelector @@ -55,9 +57,11 @@ package object dimwit: AxisAtIndex, AxisAtRange, AxisAtIndices, - AxisAtTensorIndex + AxisAtTensorIndex, + AxisAtTensorIndices, + AxisAtWindow } - export dimwit.tensor.ShapeTypeHelpers.{AxisInTensor, AxisIndex, AxisRemover, AxisReplacer, AxisIndices, AxesRemover, AxesConditionalRemover, SharedAxisRemover} + export dimwit.tensor.ShapeTypeHelpers.{AxisInTensor, AxisIndex, AxisRemover, AxisReplacer, AxisIndices, AxesRemover, SharedAxisRemover} // Export the Prime axis marker and the type classes that manipulate it export dimwit.prime.{Prime, PrimeRemover, PrimeRest, PrimeConcat} diff --git a/core/src/main/scala/dimwit/random/Random.scala b/core/src/main/scala/dimwit/random/Random.scala index d0f8ce28..12d9af93 100644 --- a/core/src/main/scala/dimwit/random/Random.scala +++ b/core/src/main/scala/dimwit/random/Random.scala @@ -73,7 +73,7 @@ object Random: * Array(3, 4), * Array(5, 6)) ) * val perm = Random.permutation(Axis[Row] -> 3)(key) - * val shuffled = data.take(Axis[Row])(perm) + * val shuffled = data.slice(Axis[Row].at(perm)) * }}} * * @param dim diff --git a/core/src/main/scala/dimwit/tensor/Axis.scala b/core/src/main/scala/dimwit/tensor/Axis.scala index 583d482e..4168b1ae 100644 --- a/core/src/main/scala/dimwit/tensor/Axis.scala +++ b/core/src/main/scala/dimwit/tensor/Axis.scala @@ -5,6 +5,8 @@ import dimwit.|*| import ShapeTypeHelpers.AxisIndex +import scala.annotation.targetName + /** Instances of this class represent an axis in a tensor with a specific label `L`. * Axis objects are used whenever an axis needs to be selected at the value level, * such as when indexing into a tensor or defining the shape of a tensor. @@ -16,6 +18,10 @@ final class Axis[L: Label]: def at(range: Range): AxisAtRange[L] = AxisAtRange(this, range) def at(indices: Seq[Int]): AxisAtIndices[L] = AxisAtIndices(this, indices) def at(index: Tensor0[Int32]): AxisAtTensorIndex[L] = AxisAtTensorIndex(this, index) + @targetName("atTensorIndices") + def at[L2](indices: Tensor1[L2, Int32]): AxisAtTensorIndices[L, L2] = AxisAtTensorIndices(this, indices) + def at(start: Int, windowSize: Int): AxisAtWindow[L] = AxisAtWindow(this, start, windowSize) + def at(start: Tensor0[Int32], windowSize: Int): AxisAtWindow[L] = AxisAtWindow(this, start, windowSize) def at[I <: NonEmptyTuple](indices: I): AxisAtTupleIndices[L, I] = AxisAtTupleIndices(this, indices) def as[U](newAxis: Axis[U]): (Axis[L], Axis[U]) = (this, newAxis) @@ -54,5 +60,15 @@ case class AxisAtIndices[L](axis: Axis[L], indices: Seq[Int]) extends AxisSelect /* Represent an axis selection by a tensor containing indices. This allows for dynamic indexing based on the contents of the tensor. */ case class AxisAtTensorIndex[L](axis: Axis[L], index: Tensor0[Int32]) extends AxisSelector[L] +/* Represent an axis selection by a vector of indices. The selected axis is replaced by the axis of the vector. */ +case class AxisAtTensorIndices[L, L2](axis: Axis[L], indices: Tensor1[L2, Int32]) extends AxisSelector[L] + +/** Represent an axis selection by a window of `windowSize` consecutive indices from `start`. + * The start may be a tensor, also a traced one (e.g. a loop index under `jit`); the size is static. + * A window with an `Int` start must fit the axis. With a tensor start, which cannot be checked, + * the start is clamped so that the window fits, as in `jax.lax.dynamic_slice`. + */ +case class AxisAtWindow[L](axis: Axis[L], start: Int | Tensor0[Int32], windowSize: Int) extends AxisSelector[L] + /* Represent an axis selection by a tuple containing indices. */ case class AxisAtTupleIndices[L, I <: NonEmptyTuple](axis: Axis[L], indices: I) extends AxisSelector[L] diff --git a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala b/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala index 59aa934b..2b15e94b 100644 --- a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala +++ b/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala @@ -164,22 +164,6 @@ object ShapeTypeHelpers: type RemainingAxes = R def indices: List[Int] = idx.indices - /** Removes [[RemovedAxis]] from a tensor shape while computing runtime indices - * for [[IndexAxes]]. - */ - trait AxesConditionalRemover[TensorShape <: Tuple, RemovedAxis <: Tuple, IndexAxes <: Tuple] extends AxesInTensor[TensorShape, IndexAxes]: - type RemainingAxes <: Tuple - - object AxesConditionalRemover: - type Aux[T <: Tuple, RA <: Tuple, IA <: Tuple, R <: Tuple] = AxesConditionalRemover[T, RA, IA] { type RemainingAxes = R } - - given bridge[T <: Tuple, RemovedAxis <: Tuple, IndexAxes <: Tuple, R <: Tuple](using - idx: AxisIndices[T, IndexAxes], - ev: RemoverAll.Aux[T, RemovedAxis, R] - ): AxesConditionalRemover.Aux[T, RemovedAxis, IndexAxes, R] = new AxesConditionalRemover[T, RemovedAxis, IndexAxes]: - type RemainingAxes = R - def indices = idx.indices - /** Removes a shared axis from multiple tensor shapes while computing runtime indices * for the remaining axes. */ diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index 1b06d924..990eca31 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -8,14 +8,14 @@ import dimwit.tensor.AxisAtIndex import dimwit.tensor.AxisAtIndices import dimwit.tensor.AxisAtRange import dimwit.tensor.AxisAtTensorIndex +import dimwit.tensor.AxisAtTensorIndices +import dimwit.tensor.AxisAtWindow import dimwit.tensor.AxisAtTupleIndices import dimwit.tensor.AxisExtent import dimwit.tensor.DType.Bool -import dimwit.tensor.DType.Int32 import dimwit.tensor.Label import dimwit.tensor.Labels import dimwit.tensor.Shape -import dimwit.tensor.ShapeTypeHelpers.AxesConditionalRemover import dimwit.tensor.ShapeTypeHelpers.AxesMerger import dimwit.tensor.ShapeTypeHelpers.AxisIndex import dimwit.tensor.ShapeTypeHelpers.AxisIndices @@ -96,46 +96,6 @@ private[dimwit] object StructuralOps: private[dimwit] object Util: - type ExtractLabel[X] = X match - case AxisAtIndex[l] => l - case AxisAtRange[l] => l - case AxisAtIndices[l] => l - case AxisAtTupleIndices[l, ?] => l - case AxisAtTensorIndex[l] => l - type ExtractLabels[Inputs <: Tuple] = Tuple.Map[Inputs, ExtractLabel] - - trait SliceLabelExtractor[Inputs <: Tuple, Out <: Tuple] - - object SliceLabelExtractor: - - given emptyTuple: SliceLabelExtractor[EmptyTuple, EmptyTuple] = - new SliceLabelExtractor[EmptyTuple, EmptyTuple] {} - - given atIndexTuple[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[AxisAtIndex[L] *: Tail, L *: TailOut] = - new SliceLabelExtractor[AxisAtIndex[L] *: Tail, L *: TailOut] {} - - given atRangeTuple[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[AxisAtRange[L] *: Tail, TailOut] = - new SliceLabelExtractor[AxisAtRange[L] *: Tail, TailOut] {} - - given atIndicesTuple[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[AxisAtIndices[L] *: Tail, TailOut] = - new SliceLabelExtractor[AxisAtIndices[L] *: Tail, TailOut] {} - - given atTupleIndicesTuple[L, I <: NonEmptyTuple, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[AxisAtTupleIndices[L, I] *: Tail, TailOut] = - new SliceLabelExtractor[AxisAtTupleIndices[L, I] *: Tail, TailOut] {} - - given atTensorIndexTuple[L, Tail <: Tuple, TailOut <: Tuple](using - tailExt: SliceLabelExtractor[Tail, TailOut] - ): SliceLabelExtractor[AxisAtTensorIndex[L] *: Tail, L *: TailOut] = - new SliceLabelExtractor[AxisAtTensorIndex[L] *: Tail, L *: TailOut] {} - @implicitNotFound("The axis ${L} is already present in the tensor shape ${T}.") trait AxisAbsent[T, L] object AxisAbsent: @@ -366,6 +326,17 @@ private[dimwit] object StructuralExtensions: import StructuralOps.* import StructuralOps.Util.* + /** The Python list of `indices`. */ + private def pyIndices(indices: Seq[Int]): py.Any = + indices.map(py.Any.from).toPythonCopy // TODO find out why Copy is needed here + + /** Requires that `value` has the shape `extents` of the selection it is written to, so that JAX does not broadcast it. */ + private def requireSelectionShape(value: Tensor[?, ?], extents: Seq[Int]): Unit = + require( + value.shape.dimensions == extents, + s"Cannot set a value of shape ${value.shape} into a selection of shape ${value.axes.zip(extents).map((a, n) => s"$a -> $n").mkString("(", ", ", ")")}" + ) + extension [T <: Tuple, V](tensor: Tensor[T, V]) /** takes a concatenated tensor and splits it into a tuple of tensors along the specified axis, @@ -535,41 +506,6 @@ private[dimwit] object StructuralExtensions: ): maker.Out = split(AxisAtTupleIndices(selector.axis, Tuple1(selector.index))) - private def calcPyIndices[Inputs <: Tuple]( - inputs: Inputs, - targetDims: List[Int] - ) = - - val PySlice = py.Dynamic.global.slice - val Colon = PySlice(py.None) - val rank = tensor.shape.rank - val indicesBuffer = collection.mutable.ArrayBuffer.fill[py.Any](rank)(Colon) - - val inputList = inputs.toList.asInstanceOf[List[Any]] - - targetDims.zip(inputList).foreach { case (dimIndex, input) => - val dimSize = tensor.shape.dimensions(dimIndex) - input match - case AxisAtIndex(_, idx) => - indicesBuffer(dimIndex) = py.Any.from(idx) - case AxisAtRange(_, range) => - require(range.isEmpty || (range.min >= 0 && range.max < dimSize), s"$range is out of bounds for axis of size $dimSize") - // map Scala Range to Python Range which is exclusive - val stop = range match - case r: Range.Inclusive => r.end + r.step.sign - case r: Range.Exclusive => r.end - // the range is in bounds, so a negative stop can only mean "before the first element", which Python spells as None - indicesBuffer(dimIndex) = PySlice(range.start, if stop < 0 then py.None else py.Any.from(stop), range.step) - case AxisAtIndices(_, indices) => - indicesBuffer(dimIndex) = indices.map(py.Any.from).toPythonCopy // TODO find out why Copy is needed here - case AxisAtTupleIndices(_, indices) => - indicesBuffer(dimIndex) = indices.toList.asInstanceOf[List[Int]].map(py.Any.from).toPythonCopy - case AxisAtTensorIndex(_, tensorIdx) => - indicesBuffer(dimIndex) = tensorIdx.jaxValue - } - - Jax.Dynamic.global.tuple(indicesBuffer.toSeq.toPythonProxy) - /** Unstacks the tensor along the specified axis at the given indices, returning a sequence of tensors corresponding to the splits. * * @param unstackAxis the axis to split, specified as an Axis (e.g. Axis[Ax1]) @@ -594,156 +530,203 @@ private[dimwit] object StructuralExtensions: val res = Jax.jnp.split(tensor.jaxValue, chunkSize, axis = axisIndex.index).as[Seq[Jax.PyDynamic]] res.map(x => Tensor[T, V](x)) - /** Slices the tensor according to the specified inputs, - * removing the specified labels from the resulting tensor. + /** Slice the tensor at an index along an axis, removing the axis. * - * @param inputs A tuple of inputs specifying how to slice the tensor. - * @return The sliced tensor with the specified labels removed from its shape. + * @param selector of the form `Axis[L].at(index)` */ - def slice[Inputs <: Tuple, LabelsToRemove <: Tuple]( - inputs: Inputs - )(using - sliceExtractor: SliceLabelExtractor[Inputs, LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Inputs]], + def slice[L](selector: AxisAtIndex[L])(using + ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes] ): Tensor[ev.RemainingAxes, V] = - val pyIndices = tensor.calcPyIndices(inputs, ev.indices) - Tensor(tensor.jaxValue.itemAt(pyIndices)) + Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index)))) - /** Slice the given tensor, specifying the axis and index to slice at. + /** Slice the tensor at an index given by a tensor along an axis, removing the axis. + * The index may be traced, e.g. a loop index under `jit`. * - * @param selector An AxisAtIndex specifying the axis and index to slice at. - * @return A sliced tensor with the specified axis removed from its shape. + * @param selector of the form `Axis[L].at(index: Tensor0[Int32])` */ - def slice[L, LabelsToRemove <: Tuple]( - selector: AxisAtIndex[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtIndex[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtIndex[L]]]], + def slice[L](selector: AxisAtTensorIndex[L])(using + ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = slice(Tuple1(selector)) + ): Tensor[ev.RemainingAxes, V] = + Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue))) - /** Slice the given tensor, specifying the axis and a given range to slice at. + /** Slice the tensor at a range along an axis, keeping the axis with the length of the range as extent. * - * @param selector An AxisAtRange specifying the axis and range to slice at. - * @return A sliced tensor with the specified axis removed from its shape. + * @param selector of the form `Axis[L].at(range)` */ - def slice[L, LabelsToRemove <: Tuple]( - selector: AxisAtRange[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtRange[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtRange[L]]]], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = slice(Tuple1(selector)) + def slice[L](selector: AxisAtRange[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + ): Tensor[T, V] = + Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index)))) - /** Slice the given tensor, specifying the axis and a list of indices to slice at. + /** Slice the tensor at a sequence of indices along an axis, keeping the axis with the number of indices as extent. * - * @param selector An AxisAtIndices specifying the axis and indices to slice at. - * @return A sliced tensor with the specified axis removed from its shape. + * @param selector of the form `Axis[L].at(Seq(i, j, ...))` */ - def slice[L, LabelsToRemove <: Tuple]( - selector: AxisAtIndices[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtIndices[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtIndices[L]]]], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = slice(Tuple1(selector)) + def slice[L](selector: AxisAtIndices[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + ): Tensor[T, V] = + Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices)))) - /** Slice the given tensor, specifying the axis and a tensor of indices to slice at. + /** Slice the tensor at a tuple of indices along an axis, keeping the axis with the number of indices as extent. * - * @param selector An AxisAtTensorIndex specifying the axis and tensor of indices to slice at. - * @return A sliced tensor with the specified axis removed from its shape. + * @param selector of the form `Axis[L].at((i, j, ...))` */ - def slice[L, LabelsToRemove <: Tuple]( - selector: AxisAtTensorIndex[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtTensorIndex[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtTensorIndex[L]]]], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = slice(Tuple1(selector)) + def slice[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + ): Tensor[T, V] = + Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices.toList.asInstanceOf[List[Int]])))) - /** Slice the given tensor, specifying the axis and a tuple of indices to slice at. + /** Slice the tensor at the indices of a vector along an axis, like `jnp.take`. + * The axis is replaced by the axis of the indices. Out-of-bounds indices yield NaN for floating tensors. * - * @param selector An AxisAtTupleIndices specifying the axis and tuple of indices to slice at. - * @return A sliced tensor with the specified axis removed from its shape. + * @param selector of the form `Axis[L].at(indices: Tensor1[L2, Int32])` */ - def slice[L, U <: NonEmptyTuple, LabelsToRemove <: Tuple]( - selector: AxisAtTupleIndices[L, U] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtTupleIndices[L, U]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtTupleIndices[L, U]]]], - labels: Labels[ev.RemainingAxes] - ): Tensor[ev.RemainingAxes, V] = slice(Tuple1(selector)) - - def take[L1, L2: Label]( - axis: Axis[L1] - )( - indices: Tensor1[L2, Int32] - )(using - ev: AxisReplacer[T, L1, L2], + def slice[L, L2](selector: AxisAtTensorIndices[L, L2])(using + ev: AxisReplacer[T, L, L2], labels: Labels[ev.NewShape] ): Tensor[ev.NewShape, V] = - val result = Jax.jnp.take(tensor.jaxValue, indices.jaxValue, axis = ev.index) - Tensor(result) + Tensor(Jax.jnp.take(tensor.jaxValue, selector.indices.jaxValue, axis = ev.index)) - def set[Inputs <: Tuple, LabelsToRemove <: Tuple]( - inputs: Inputs - )(using - sliceExtractor: SliceLabelExtractor[Inputs, LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Inputs]], + /** Slice a window of static size along an axis, like `jax.lax.dynamic_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 window, which keeps the axis with extent `windowSize`. + */ + def slice[L](selector: AxisAtWindow[L])(using + axisIndex: AxisIndex[T, L], + labels: Labels[T] + ): Tensor[T, V] = + val window = tensor.windowExtents(selector, axisIndex.index) + Tensor(Jax.lax.dynamic_slice(tensor.jaxValue, tensor.windowStarts(selector, axisIndex.index), window.toPythonCopy)) + + /** Overwrite 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.slice(selector)` + */ + def set[L](selector: AxisAtIndex[L])(using + ev: AxisRemover[T, L], labels: Labels[T] )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = - val pyIndices = tensor.calcPyIndices(inputs, ev.indices) - val result = tensor.jaxValue.at.itemAt(pyIndices).set(value.jaxValue) - Tensor(result) + 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)) - // Convenience overload for Float - def set[Inputs <: Tuple, LabelsToRemove <: Tuple]( - inputs: Inputs - )(using - sliceExtractor: SliceLabelExtractor[Inputs, LabelsToRemove], - ev: AxesConditionalRemover.Aux[T, LabelsToRemove, ExtractLabels[Inputs], EmptyTuple], + /** Overwrite 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.slice(selector)` + */ + def set[L](selector: AxisAtTensorIndex[L])(using + ev: AxisRemover[T, L], labels: Labels[T] - )(value: Float): Tensor[T, V] = - val pyIndices = tensor.calcPyIndices(inputs, ev.indices) - val result = tensor.jaxValue.at.itemAt(pyIndices).set(value) - Tensor(result) - - // Convenience overload for AxisAtIndex - def set[L, LabelsToRemove <: Tuple]( - selector: AxisAtIndex[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtIndex[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtIndex[L]]]], + )(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)) + + /** Overwrite 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.slice(selector)` + */ + def set[L](selector: AxisAtRange[L])(using + axisIndex: AxisIndex[T, L], labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = set(Tuple1(selector))(value) + )(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)) - // Convenience overload for AxisAtRange - def set[L, LabelsToRemove <: Tuple]( - selector: AxisAtRange[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtRange[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtRange[L]]]], + /** Overwrite 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.slice(selector)` + */ + def set[L](selector: AxisAtIndices[L])(using + axisIndex: AxisIndex[T, L], labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = set(Tuple1(selector))(value) + )(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)) - // Convenience overload for AxisAtIndices - def set[L, LabelsToRemove <: Tuple]( - selector: AxisAtIndices[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtIndices[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtIndices[L]]]], + /** Overwrite 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.slice(selector)` + */ + def set[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using + axisIndex: AxisIndex[T, L], labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = set(Tuple1(selector))(value) + )(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)) - // Convenience overload for AxisAtTensorIndex - def set[L, LabelsToRemove <: Tuple]( - selector: AxisAtTensorIndex[L] - )(using - sliceExtractor: SliceLabelExtractor[Tuple1[AxisAtTensorIndex[L]], LabelsToRemove], - ev: AxesConditionalRemover[T, LabelsToRemove, ExtractLabels[Tuple1[AxisAtTensorIndex[L]]]], + /** Overwrite 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.slice(selector)` + */ + def set[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)) + + /** Overwrite 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.slice(selector)` + */ + def set[L](selector: AxisAtWindow[L])(using + axisIndex: AxisIndex[T, L], labels: Labels[T] - )(value: Tensor[ev.RemainingAxes, V]): Tensor[T, V] = set(Tuple1(selector))(value) + )(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))) + + /** 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) + Jax.Dynamic.global.tuple(Seq.tabulate[py.Any](tensor.shape.rank)(i => if i == dimIndex then index else all).toPythonProxy) + + /** The Python slice for `range` on the axis `dimIndex`, which the range must not leave. */ + private def pyRange(range: Range, dimIndex: Int): py.Any = + val dimSize = tensor.shape.dimensions(dimIndex) + require(range.isEmpty || (range.min >= 0 && range.max < dimSize), s"$range is out of bounds for axis of size $dimSize") + // map Scala Range to Python Range which is exclusive + val stop = range match + case r: Range.Inclusive => r.end + r.step.sign + case r: Range.Exclusive => r.end + // the range is in bounds, so a negative stop can only mean "before the first element", which Python spells as None + py.Dynamic.global.slice(range.start, if stop < 0 then py.None else py.Any.from(stop), range.step) + + /** The extents of the window `selector` on the axis at `dimIndex`: `windowSize` there, the whole axis elsewhere. */ + private def windowExtents(selector: AxisAtWindow[?], dimIndex: Int): Seq[Int] = + val dimSize = tensor.shape.dimensions(dimIndex) + val size = selector.windowSize + require(size >= 0 && size <= dimSize, s"Window of size $size does not fit axis ${tensor.axes(dimIndex)} of size $dimSize") + selector.start match + case start: Int => + require(start >= 0 && start + size <= dimSize, s"Window $start until ${start + size} is out of bounds for axis ${tensor.axes(dimIndex)} of size $dimSize") + case _ => () // a tensor start may be traced, so JAX clamps it + tensor.shape.dimensions.updated(dimIndex, size) + + /** The start of the window `selector` on every axis: its start on the axis at `dimIndex`, 0 elsewhere. */ + private def windowStarts(selector: AxisAtWindow[?], dimIndex: Int): Jax.PyDynamic = + val start: py.Any = selector.start match + case start: Int => py.Any.from(start) + case start: Tensor[?, ?] => start.jaxValue + val starts = Seq.tabulate[py.Any](tensor.shape.rank)(i => if i == dimIndex then start else py.Any.from(0)) + Jax.Dynamic.global.tuple(starts.toPythonProxy) def rearrange[Axes <: Tuple, Status <: ValidationResult](newOrder: Axes)(using Labels[UnwrapAxes[Axes]] diff --git a/core/src/main/scala/dimwit/tensor/tensorops/Tensor1Extensions.scala b/core/src/main/scala/dimwit/tensor/tensorops/Tensor1Extensions.scala index 95114bd5..32434958 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/Tensor1Extensions.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/Tensor1Extensions.scala @@ -1,17 +1,10 @@ package dimwit.tensor.tensorops -import dimwit.jax.Jax import dimwit.tensor.Axis -import dimwit.tensor.DType.* import dimwit.tensor.HasScalar import dimwit.tensor.Label -import dimwit.tensor.Labels import dimwit.tensor.Tensor -import dimwit.tensor.Tensor0 import dimwit.tensor.Tensor1 -import me.shadaj.scalapy.py -import me.shadaj.scalapy.py.SeqConverters -import me.shadaj.scalapy.readwrite.Writer private[dimwit] object Tensor1Extensions: @@ -23,16 +16,6 @@ private[dimwit] object Tensor1Extensions: */ def relabelTo[NewL: Label](newAxis: Axis[NewL]): Tensor1[NewL, V] = Tensor[Tuple1[NewL], V](t.jaxValue) - // TODO generalize to TensorN (like slice) - def dynamicSlice( - dynamicStart: Tensor0[Int32], - staticSize: Int - )(using - label: Label[L] - ): Tensor1[L, V] = - // TODO understand why toPythonCopy is needed and toPythonProxy fails! - Tensor(Jax.lax.dynamic_slice(t.jaxValue, Seq(dynamicStart.jaxValue).toPythonCopy, Seq(staticSize).toPythonCopy)) - extension [L, V, X](t: Tensor1[L, V])(using ev: HasScalar[V, X]) /** Converts a Tensor1 to a Scala Array. * The user must ensure that the tensor is not a JAX Tracer diff --git a/core/src/test/scala/dimwit/random/RandomSuite.scala b/core/src/test/scala/dimwit/random/RandomSuite.scala index 813bda2a..113fa724 100644 --- a/core/src/test/scala/dimwit/random/RandomSuite.scala +++ b/core/src/test/scala/dimwit/random/RandomSuite.scala @@ -99,7 +99,7 @@ class RandomSuite extends DimwitTest: )) val rowPerm = Random.permutation(Axis[Row] -> 4)(key) - val shuffled = original.take(Axis[Row])(rowPerm) + val shuffled = original.slice(Axis[Row].at(rowPerm)) shuffled.shape should equal(original.shape) shuffled.sum.item shouldBe original.sum.item // sum should be unchanged diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala index 5d6a63e8..3d3f843b 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -164,9 +164,9 @@ class TensorCreationSuite extends DimwitTest: typeCheckErrors("Tensor1(Axis[A]).fromRange(0 until 3, VType[Float32])") should not be empty typeCheckErrors("Tensor1(Axis[A], VType[Float32]).fromRange(0 until 3)") should not be empty - it("can be consumed as gather indices by take"): + it("can be consumed as gather indices by slice"): val t = Tensor1(Axis[A]).fromArray(Array(10.0f, 20.0f, 30.0f)) - t.take(Axis[A])(Tensor1(Axis[B]).fromRange(0 until 3)) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) + t.slice(Axis[A].at(Tensor1(Axis[B]).fromRange(0 until 3))) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) describe("linspace"): diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala index f50f99a8..ed82688f 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala @@ -162,10 +162,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = AB.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice((Axis[C].at(0), Axis[D].at(0))) should approxEqual(AB) - res.slice((Axis[C].at(1), Axis[D].at(0))) should approxEqual(AB) - res.slice((Axis[C].at(0), Axis[D].at(1))) should approxEqual(AB) - res.slice((Axis[C].at(1), Axis[D].at(1))) should approxEqual(AB) + res.slice(Axis[C].at(0)).slice(Axis[D].at(0)) should approxEqual(AB) + res.slice(Axis[C].at(1)).slice(Axis[D].at(0)) should approxEqual(AB) + res.slice(Axis[C].at(0)).slice(Axis[D].at(1)) should approxEqual(AB) + res.slice(Axis[C].at(1)).slice(Axis[D].at(1)) should approxEqual(AB) it("BC broadcastTo ABCD"): val BC = Tensor(Shape(Axis[B] -> 2, Axis[C] -> 2)).fromArray( @@ -173,10 +173,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = BC.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice((Axis[A].at(0), Axis[D].at(0))) should approxEqual(BC) - res.slice((Axis[A].at(1), Axis[D].at(0))) should approxEqual(BC) - res.slice((Axis[A].at(0), Axis[D].at(1))) should approxEqual(BC) - res.slice((Axis[A].at(1), Axis[D].at(1))) should approxEqual(BC) + res.slice(Axis[A].at(0)).slice(Axis[D].at(0)) should approxEqual(BC) + res.slice(Axis[A].at(1)).slice(Axis[D].at(0)) should approxEqual(BC) + res.slice(Axis[A].at(0)).slice(Axis[D].at(1)) should approxEqual(BC) + res.slice(Axis[A].at(1)).slice(Axis[D].at(1)) should approxEqual(BC) it("CD broadcastTo ABCD"): val CD = Tensor(Shape(Axis[C] -> 2, Axis[D] -> 2)).fromArray( @@ -184,10 +184,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = CD.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice((Axis[A].at(0), Axis[B].at(0))) should approxEqual(CD) - res.slice((Axis[A].at(1), Axis[B].at(0))) should approxEqual(CD) - res.slice((Axis[A].at(0), Axis[B].at(1))) should approxEqual(CD) - res.slice((Axis[A].at(1), Axis[B].at(1))) should approxEqual(CD) + res.slice(Axis[A].at(0)).slice(Axis[B].at(0)) should approxEqual(CD) + res.slice(Axis[A].at(1)).slice(Axis[B].at(0)) should approxEqual(CD) + res.slice(Axis[A].at(0)).slice(Axis[B].at(1)) should approxEqual(CD) + res.slice(Axis[A].at(1)).slice(Axis[B].at(1)) should approxEqual(CD) describe("Disallow"): diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index f9f0c4a3..0485dbcf 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -490,18 +490,29 @@ class TensorOpsStructureSuite extends DimwitTest: s2 should approxEqual(part2) s3 should approxEqual(part3) + describe("selectors"): + + it("are applied one at a time, not as a tuple"): + "t3.slice((Axis[A].at(0), Axis[B].at(0)))" shouldNot compile + "t3.set((Axis[A].at(0), Axis[B].at(0)))(Tensor1(Axis[C]).fromArray(Array(0f)))" shouldNot compile + + it("reject a value whose shape differs from the selection"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) + the[IllegalArgumentException] thrownBy t.set(Axis[A].at(0 to 1))(t) should have message + "requirement failed: Cannot set a value of shape Shape(A -> 3) into a selection of shape (A -> 2)" + describe("set function"): describe("AxisAtIndex"): it("set single value (scalar index)"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - t = t.set((Axis[A].at(0), Axis[B].at(0)))(5f) + t = t.set(Axis[A].at(0))(t.slice(Axis[A].at(0)).set(Axis[B].at(0))(5f)) t should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(5.0f, 2.0f), Array(3.0f, 4.0f)))) it("set single value (tensor index)"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - t = t.set((Axis[A].at(Tensor0(1)), Axis[B].at(Tensor0(1))))(5f) + t = t.set(Axis[A].at(Tensor0(1)))(t.slice(Axis[A].at(Tensor0(1))).set(Axis[B].at(Tensor0(1)))(5f)) t should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 5.0f)))) it("set a sub-vector"): @@ -567,6 +578,31 @@ class TensorOpsStructureSuite extends DimwitTest: val t2 = t.set(Axis[A].at(List(0, 2)))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) t2 should approxEqual(Tensor1(Axis[A]).fromArray(Array(7.0f, 2.0f, 8.0f))) + describe("AxisAtTensorIndices"): + + // Shape: A=2, B=3, the value at (a, b) is 10a + b + val ab = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(2, 3)((a, b) => 10f * a + b)) + val indices = Tensor1(Axis[C]).fromArray(Array(2, 0, 2, 1)) + + it("slices at the indices and replaces the axis by the axis of the indices"): + val sliced = ab.slice(Axis[B].at(indices)) + sliced.axes shouldBe List("A", "C") + sliced should approxEqual(Tensor2(Axis[A], Axis[C]).fromArray(Array(Array(2f, 0f, 2f, 1f), Array(12f, 10f, 12f, 11f)))) + + it("keeps the axis when the indices have the same axis"): + val permutation = Tensor1(Axis[B]).fromArray(Array(2, 0, 1)) + ab.slice(Axis[B].at(permutation)) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(12f, 10f, 11f)))) + + it("sets at the indices"): + val zeros = Tensor(ab.shape).fill(0f) + val positions = Tensor1(Axis[C]).fromArray(Array(2, 0)) + val values = Tensor2(Axis[A], Axis[C]).fromArray(Array(Array(1f, 2f), Array(3f, 4f))) + zeros.set(Axis[B].at(positions))(values) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(4f, 0f, 3f)))) + + it("set is the inverse of slice for distinct indices"): + val permutation = Tensor1(Axis[B]).fromArray(Array(2, 0, 1)) + Tensor(ab.shape).fill(0f).set(Axis[B].at(permutation))(ab.slice(Axis[B].at(permutation))) should approxEqual(ab) + describe("roll"): it("tensor1"): val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) @@ -579,6 +615,96 @@ 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("window (AxisAtWindow)"): + + // Shape: A=2, B=5, C=3, the value at (a, b, c) is 100a + 10b + c + val abc = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( + Array.tabulate(2, 5, 3)((a, b, c) => 100f * a + 10f * b + c) + ) + + describe("slice"): + + it("slices a window of a vector"): + val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f)) + t.slice(Axis[A].at(Tensor0(2), 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f, 5.0f))) + + it("slices a window along a middle axis and keeps the axis"): + val window: Tensor3[A, B, C, Float32] = abc.slice(Axis[B].at(Tensor0(1), 2)) + window.axes shouldBe List("A", "B", "C") + window.shape.dimensions shouldBe List(2, 2, 3) + window should approxEqual(abc.slice(Axis[B].at(1 until 3))) + + it("clamps the start so that the window fits"): + abc.slice(Axis[B].at(Tensor0(4), 3)) should approxEqual(abc.slice(Axis[B].at(2 until 5))) + + it("chains with other selectors"): + val window = abc.slice(Axis[B].at(Tensor0(1), 2)).slice(Axis[C].at(0)) + window.axes shouldBe List("A", "B") + window should approxEqual(abc.slice(Axis[B].at(1 until 3)).slice(Axis[C].at(0))) + + it("rejects an axis that the tensor does not have"): + "abc.slice(Axis[D].at(Tensor0(1), 2))" shouldNot compile + + it("works with a traced start under jit"): + val sliceAt = jit((start: Tensor0[Int32]) => abc.slice(Axis[B].at(start, 2))) + sliceAt(Tensor0(0)) should approxEqual(abc.slice(Axis[B].at(0 until 2))) + sliceAt(Tensor0(3)) should approxEqual(abc.slice(Axis[B].at(3 until 5))) + + it("is differentiable"): + val grad = Autodiff.grad((t: Tensor3[A, B, C, Float32]) => t.slice(Axis[B].at(Tensor0(1), 2)).sum) + val expected = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( + Array.tabulate(2, 5, 3)((_, b, _) => if b == 1 || b == 2 then 1f else 0f) + ) + grad(abc).value should approxEqual(expected) + + it("slices a window with an Int start"): + abc.slice(Axis[B].at(1, 2)) should approxEqual(abc.slice(Axis[B].at(1 until 3))) + + it("rejects a window with an Int start that does not fit"): + the[IllegalArgumentException] thrownBy abc.slice(Axis[B].at(4, 3)) should have message + "requirement failed: Window 4 until 7 is out of bounds for axis B of size 5" + + it("rejects a window larger than the axis"): + the[IllegalArgumentException] thrownBy abc.slice(Axis[B].at(Tensor0(0), 6)) should have message + "requirement failed: Window of size 6 does not fit axis B of size 5" + + describe("set"): + + val zeros = Tensor(abc.shape).fill(0f) + val block = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( + Array.tabulate(2, 2, 3)((a, b, c) => 1f + 100f * a + 10f * b + c) + ) + + it("sets a window along a middle axis"): + zeros.set(Axis[B].at(Tensor0(1), 2))(block) should approxEqual(zeros.set(Axis[B].at(1 until 3))(block)) + + it("clamps the start so that the window fits"): + zeros.set(Axis[B].at(Tensor0(4), 2))(block) should approxEqual(zeros.set(Axis[B].at(3 until 5))(block)) + + it("works with a traced start under jit"): + val setAt = jit((start: Tensor0[Int32]) => zeros.set(Axis[B].at(start, 2))(block)) + setAt(Tensor0(0)) should approxEqual(zeros.set(Axis[B].at(0 until 2))(block)) + setAt(Tensor0(3)) should approxEqual(zeros.set(Axis[B].at(3 until 5))(block)) + + it("sets a window with an Int start"): + zeros.set(Axis[B].at(1, 2))(block) should approxEqual(zeros.set(Axis[B].at(1 until 3))(block)) + + it("fills a buffer row by row"): + val filled = (0 until 5).foldLeft(zeros): (buffer, b) => + buffer.set(Axis[B].at(Tensor0(b), 1))(abc.slice(Axis[B].at(b until b + 1))) + filled should approxEqual(abc) + + it("is differentiable with respect to the value"): + val grad = Autodiff.grad((value: Tensor3[A, B, C, Float32]) => (zeros.set(Axis[B].at(Tensor0(1), 2))(value) * abc).sum) + grad(block).value should approxEqual(abc.slice(Axis[B].at(1 until 3))) + + it("rejects a value whose extent differs from the window size"): + the[IllegalArgumentException] thrownBy zeros.set(Axis[B].at(Tensor0(1), 3))(block) should have message + "requirement failed: Cannot set a value of shape Shape(A -> 2, B -> 2, C -> 3) into a selection of shape (A -> 2, B -> 3, C -> 3)" + + it("does not broadcast a value without the window axis"): + "zeros.set(Axis[B].at(Tensor0(1), 2))(block.slice(Axis[B].at(0)))" shouldNot compile + describe("chunk function"): it("chunk axis into equal parts"): diff --git a/examples/src/main/scala/dimwit/basic/KMeans.scala b/examples/src/main/scala/dimwit/basic/KMeans.scala index 9a594c51..9719a2f0 100644 --- a/examples/src/main/scala/dimwit/basic/KMeans.scala +++ b/examples/src/main/scala/dimwit/basic/KMeans.scala @@ -114,7 +114,7 @@ trait CenterBasedClustering( key: Random.Key ): Tensor2[Cluster, Dim, Float32] = val perm = Random.permutation(points.shape.extent(Axis[Point]))(key) - points.take(Axis[Point])(perm.slice(Axis[Point].at(0 until nClusters))).relabel(Axis[Point].as(Axis[Cluster])) + points.slice(Axis[Point].at(perm.slice(Axis[Point].at(0 until nClusters)))).relabel(Axis[Point].as(Axis[Cluster])) /** Example implementation of the K-Means clustering algorithm, showcasing the use * of Named Tensors for clear and type-safe code. diff --git a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala index 12eee4d3..471e60de 100644 --- a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala +++ b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala @@ -90,10 +90,10 @@ object LogisticRegression: val testPerm = perm.slice(Axis[Sample].at(splitIndex until numSamples)) // Use the permutations to get our training and validation data - val trainingDataUnnormalized = dataInitial.take(Axis[Sample])(trainPerm) - val valDataUnnormalized = dataInitial.take(Axis[Sample])(testPerm) - val trainLabels = labelsInitial.take(Axis[Sample])(trainPerm) - val valLabels = labelsInitial.take(Axis[Sample])(testPerm) + val trainingDataUnnormalized = dataInitial.slice(Axis[Sample].at(trainPerm)) + val valDataUnnormalized = dataInitial.slice(Axis[Sample].at(testPerm)) + val trainLabels = labelsInitial.slice(Axis[Sample].at(trainPerm)) + val valLabels = labelsInitial.slice(Axis[Sample].at(testPerm)) def calcMeanAndStd(t: Tensor2[Sample, Feature, Float32]): (Tensor1[Feature, Float32], Tensor1[Feature, Float32]) = val mean = t.vmap(Axis[Feature])(_.mean) From 19776520a5a37af2c3ec098e8ad29f169ff4944f Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Wed, 7 Oct 2026 18:11:20 +0200 Subject: [PATCH 2/3] Add slice operation to AGENTS.md --- AGENTS.md | 94 +++++++++++++++++++++++++++++++++++++++++++++++-- mdocs/AGENTS.md | 63 +++++++++++++++++++++++++++++++-- 2 files changed, 153 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 1ed0002d..b3b7bb6c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -647,11 +647,11 @@ println(s"Transposed shape: ${transposed.shape}") val vec1d = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) val vec2d = vec1d.appendAxis(Axis[B]) // Add new axis B -// Take (indexing with single index) +// Slice at an index (removes the axis) val firstRow = original.slice(Axis[A].at(0)) val secondRow = original.slice(Axis[A].at(1)) -// Take (indexing with range) +// Slice at a range (keeps the axis) val data3d = Tensor(Shape3(Axis[A] -> 5, Axis[B] -> 3, Axis[C] -> 4)).fill(1.0f) val middleSlice = data3d.slice(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 println(s"Middle slice shape: ${middleSlice.shape}") // Shape(A -> 3, B -> 3, C -> 4) @@ -662,6 +662,96 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) val concatenated = concatenate(t1, t2, Axis[A]) ``` +### Slicing and Setting + +Every selection along an axis is an `Axis[L].at(...)` selector. `slice` reads the selection, +`set` returns a copy with the selection overwritten. A selector selects along exactly one axis; +to select along several axes, chain the calls. + +```scala +// Shape: A=4, B=3, the value at (a, b) is 10a + b +val grid = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(4, 3)((a, b) => 10.0f * a + b)) + +// Index: removes the axis +val gridRow: Tensor1[B, Float32] = grid.slice(Axis[A].at(1)) // [10, 11, 12] + +// Range, sequence or tuple of indices: keeps the axis, with the number of indices as extent +val gridRows = grid.slice(Axis[A].at(1 until 3)) // rows 1, 2 +val gridPicked = grid.slice(Axis[A].at(Seq(3, 0))) // rows 3, 0 +val gridTupled = grid.slice(Axis[A].at((3, 0))) // the same, note the double parentheses + +// Vector of indices, like jnp.take: the axis is replaced by the axis of the indices +val order = Tensor1(Axis[C]).fromArray(Array(2, 0, 2)) +val gridGathered: Tensor2[A, C, Float32] = grid.slice(Axis[B].at(order)) + +// Window of static size, like jax.lax.dynamic_slice: keeps the axis with the window size as extent +val gridWindow = grid.slice(Axis[A].at(1, 2)) // rows 1, 2 + +// Several axes: chain the calls +val gridEntry: Tensor0[Float32] = grid.slice(Axis[A].at(1)).slice(Axis[B].at(2)) // 12 + +// set takes a value with the shape of the selection +val gridZeroRow = grid.set(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) +val gridZeroWindow = grid.set(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) +``` + +The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one, +e.g. a loop position under `jit`. Its value is then only known at runtime, but the window size stays static: + +```scala +// Write a row of ones into a buffer at a position known only at runtime +val writeRow = jit((buffer: Tensor2[A, B, Float32], position: Tensor0[Int32]) => + buffer.set(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) +) +val gridWritten = writeRow(grid, Tensor0(2)) // row 2 is now ones +``` + +```scala +// ERROR: a selector selects along one axis, chain the calls instead +val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) +// error: +// value slice is not a member of dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, +// dimwit.tensor.DType.Float32]. +// An extension method was tried, but could not be fully constructed: +// +// dimwit.slice() +// +// failed with: +// +// value slice: does not take parameters +// val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) +// ^^^^^^^^^^ +``` + +```scala +// ERROR: an index removes the axis, so the value must not have it +val gridWrongValue = grid.set(Axis[A].at(0))(grid) +// error: +// Found: (MdocApp1.this.grid : +// dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, +// dimwit.tensor.DType.Float32] +// ) +// Required: dimwit.tensor.Tensor[?1.RemainingAxes, dimwit.tensor.DType.Float32] +// +// where: ?1 is an unknown value of type dimwit.tensor.ShapeTypeHelpers.AxisRemover.Aux[ +// (MdocApp1.this.A, MdocApp1.this.B), MdocApp1.this.A, MdocApp1.this.B *: +// EmptyTuple.type] +// +// val gridWrongValue = grid.set(Axis[A].at(0))(grid) +// ^^^^ +``` + +```scala +// ERROR at runtime: a window with an Int start must lie inside the axis +grid.slice(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:719) +// at dimwit.tensor.tensorops.StructuralExtensions$.slice(StructuralOps.scala:605) +// at dimwit.package$.slice(package.scala:80) +// at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:719) +``` + --- ## Functional Patterns diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index a0452955..f9b97afd 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -467,11 +467,11 @@ println(s"Transposed shape: ${transposed.shape}") val vec1d = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) val vec2d = vec1d.appendAxis(Axis[B]) // Add new axis B -// Take (indexing with single index) +// Slice at an index (removes the axis) val firstRow = original.slice(Axis[A].at(0)) val secondRow = original.slice(Axis[A].at(1)) -// Take (indexing with range) +// Slice at a range (keeps the axis) val data3d = Tensor(Shape3(Axis[A] -> 5, Axis[B] -> 3, Axis[C] -> 4)).fill(1.0f) val middleSlice = data3d.slice(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 println(s"Middle slice shape: ${middleSlice.shape}") // Shape(A -> 3, B -> 3, C -> 4) @@ -482,6 +482,65 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) val concatenated = concatenate(t1, t2, Axis[A]) ``` +### Slicing and Setting + +Every selection along an axis is an `Axis[L].at(...)` selector. `slice` reads the selection, +`set` returns a copy with the selection overwritten. A selector selects along exactly one axis; +to select along several axes, chain the calls. + +```scala mdoc:silent +// Shape: A=4, B=3, the value at (a, b) is 10a + b +val grid = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(4, 3)((a, b) => 10.0f * a + b)) + +// Index: removes the axis +val gridRow: Tensor1[B, Float32] = grid.slice(Axis[A].at(1)) // [10, 11, 12] + +// Range, sequence or tuple of indices: keeps the axis, with the number of indices as extent +val gridRows = grid.slice(Axis[A].at(1 until 3)) // rows 1, 2 +val gridPicked = grid.slice(Axis[A].at(Seq(3, 0))) // rows 3, 0 +val gridTupled = grid.slice(Axis[A].at((3, 0))) // the same, note the double parentheses + +// Vector of indices, like jnp.take: the axis is replaced by the axis of the indices +val order = Tensor1(Axis[C]).fromArray(Array(2, 0, 2)) +val gridGathered: Tensor2[A, C, Float32] = grid.slice(Axis[B].at(order)) + +// Window of static size, like jax.lax.dynamic_slice: keeps the axis with the window size as extent +val gridWindow = grid.slice(Axis[A].at(1, 2)) // rows 1, 2 + +// Several axes: chain the calls +val gridEntry: Tensor0[Float32] = grid.slice(Axis[A].at(1)).slice(Axis[B].at(2)) // 12 + +// set takes a value with the shape of the selection +val gridZeroRow = grid.set(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) +val gridZeroWindow = grid.set(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) +``` + +The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one, +e.g. a loop position under `jit`. Its value is then only known at runtime, but the window size stays static: + +```scala mdoc:silent +// Write a row of ones into a buffer at a position known only at runtime +val writeRow = jit((buffer: Tensor2[A, B, Float32], position: Tensor0[Int32]) => + buffer.set(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) +) +val gridWritten = writeRow(grid, Tensor0(2)) // row 2 is now ones +``` + +```scala mdoc:fail +// ERROR: a selector selects along one axis, chain the calls instead +val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) +``` + +```scala mdoc:fail +// ERROR: an index removes the axis, so the value must not have it +val gridWrongValue = grid.set(Axis[A].at(0))(grid) +``` + +```scala mdoc:crash +// ERROR at runtime: a window with an Int start must lie inside the axis +grid.slice(Axis[A].at(3, 2)) +``` + --- ## Functional Patterns From 092ec891629ab12ff127a3bfcb2c4b3abe901084 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Thu, 22 Oct 2026 11:16:03 +0200 Subject: [PATCH 3/3] Rename slice to select and set to updated, add apply on tensors `tensor(Axis[L].at(...))` forwards to select, so tensors read like Scala collections (`apply`/`updated`). Updates tests, examples and docs. --- AGENTS.md | 74 ++++---- .../src/main/scala/dimwit/random/Random.scala | 2 +- .../tensor/tensorops/StructuralOps.scala | 95 ++++++---- .../tensor/tensorops/Tensor0Extensions.scala | 2 +- .../dimwit/linalg/LinearAlgebraSuite.scala | 2 +- .../scala/dimwit/random/RandomSuite.scala | 6 +- .../dimwit/tensor/TensorCreationSuite.scala | 4 +- .../tensor/TensorOpsBroadcastSuite.scala | 24 +-- .../tensor/TensorOpsStructureSuite.scala | 174 ++++++++++-------- docs/quickstart.md | 27 ++- .../src/main/scala/dimwit/basic/KMeans.scala | 10 +- .../dimwit/basic/LogisticRegression.scala | 12 +- .../scala/dimwit/basic/SIRSimulation.scala | 8 +- .../complex/VariationalAutoencoder.scala | 3 +- mdocs/AGENTS.md | 52 +++--- mdocs/docs/quickstart.md | 27 ++- 16 files changed, 295 insertions(+), 227 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b3b7bb6c..ec79c480 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -647,14 +647,14 @@ println(s"Transposed shape: ${transposed.shape}") val vec1d = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) val vec2d = vec1d.appendAxis(Axis[B]) // Add new axis B -// Slice at an index (removes the axis) -val firstRow = original.slice(Axis[A].at(0)) -val secondRow = original.slice(Axis[A].at(1)) +// Select at an index (removes the axis) +val firstRow = original.select(Axis[A].at(0)) +val secondRow = original.select(Axis[A].at(1)) -// Slice at a range (keeps the axis) +// Select at a range (keeps the axis) val data3d = Tensor(Shape3(Axis[A] -> 5, Axis[B] -> 3, Axis[C] -> 4)).fill(1.0f) -val middleSlice = data3d.slice(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 -println(s"Middle slice shape: ${middleSlice.shape}") // Shape(A -> 3, B -> 3, C -> 4) +val middleSelection = data3d.select(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 +println(s"Middle selection shape: ${middleSelection.shape}") // Shape(A -> 3, B -> 3, C -> 4) // Concatenate val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) @@ -662,37 +662,41 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) val concatenated = concatenate(t1, t2, Axis[A]) ``` -### Slicing and Setting +### Selecting and Updating -Every selection along an axis is an `Axis[L].at(...)` selector. `slice` reads the selection, -`set` returns a copy with the selection overwritten. A selector selects along exactly one axis; -to select along several axes, chain the calls. +Every selection along an axis is an `Axis[L].at(...)` selector. `select` reads the selection, and applying +the tensor to the selector, `t(selector)`, is the same as `t.select(selector)`, as `seq(i)` reads a Scala collection. +`updated` returns a copy with new values at the selection, as `updated` does for Scala's immutable collections. +A selector selects along exactly one axis; to select along several axes, chain the calls. ```scala // Shape: A=4, B=3, the value at (a, b) is 10a + b val grid = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(4, 3)((a, b) => 10.0f * a + b)) // Index: removes the axis -val gridRow: Tensor1[B, Float32] = grid.slice(Axis[A].at(1)) // [10, 11, 12] +val gridRow: Tensor1[B, Float32] = grid.select(Axis[A].at(1)) // [10, 11, 12] // Range, sequence or tuple of indices: keeps the axis, with the number of indices as extent -val gridRows = grid.slice(Axis[A].at(1 until 3)) // rows 1, 2 -val gridPicked = grid.slice(Axis[A].at(Seq(3, 0))) // rows 3, 0 -val gridTupled = grid.slice(Axis[A].at((3, 0))) // the same, note the double parentheses +val gridRows = grid.select(Axis[A].at(1 until 3)) // rows 1, 2 +val gridPicked = grid.select(Axis[A].at(Seq(3, 0))) // rows 3, 0 +val gridTupled = grid.select(Axis[A].at((3, 0))) // the same, note the double parentheses // Vector of indices, like jnp.take: the axis is replaced by the axis of the indices val order = Tensor1(Axis[C]).fromArray(Array(2, 0, 2)) -val gridGathered: Tensor2[A, C, Float32] = grid.slice(Axis[B].at(order)) +val gridGathered: Tensor2[A, C, Float32] = grid.select(Axis[B].at(order)) // Window of static size, like jax.lax.dynamic_slice: keeps the axis with the window size as extent -val gridWindow = grid.slice(Axis[A].at(1, 2)) // rows 1, 2 +val gridWindow = grid.select(Axis[A].at(1, 2)) // rows 1, 2 // Several axes: chain the calls -val gridEntry: Tensor0[Float32] = grid.slice(Axis[A].at(1)).slice(Axis[B].at(2)) // 12 +val gridEntry: Tensor0[Float32] = grid.select(Axis[A].at(1)).select(Axis[B].at(2)) // 12 -// set takes a value with the shape of the selection -val gridZeroRow = grid.set(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) -val gridZeroWindow = grid.set(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) +// Apply: the same as select +val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 12 + +// updated takes a value with the shape of the selection +val gridZeroRow = grid.updated(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) +val gridZeroWindow = grid.updated(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) ``` The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one, @@ -701,31 +705,31 @@ e.g. a loop position under `jit`. Its value is then only known at runtime, but t ```scala // Write a row of ones into a buffer at a position known only at runtime val writeRow = jit((buffer: Tensor2[A, B, Float32], position: Tensor0[Int32]) => - buffer.set(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) + buffer.updated(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) ) val gridWritten = writeRow(grid, Tensor0(2)) // row 2 is now ones ``` ```scala // ERROR: a selector selects along one axis, chain the calls instead -val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) +val gridTuple = grid.select((Axis[A].at(1), Axis[B].at(2))) // error: -// value slice is not a member of dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, +// value select is not a member of dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, // dimwit.tensor.DType.Float32]. // An extension method was tried, but could not be fully constructed: // -// dimwit.slice() +// dimwit.select() // // failed with: // -// value slice: does not take parameters -// val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) -// ^^^^^^^^^^ +// value select: does not take parameters +// val gridTuple = grid.select((Axis[A].at(1), Axis[B].at(2))) +// ^^^^^^^^^^^ ``` ```scala // ERROR: an index removes the axis, so the value must not have it -val gridWrongValue = grid.set(Axis[A].at(0))(grid) +val gridWrongValue = grid.updated(Axis[A].at(0))(grid) // error: // Found: (MdocApp1.this.grid : // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, @@ -737,19 +741,19 @@ val gridWrongValue = grid.set(Axis[A].at(0))(grid) // (MdocApp1.this.A, MdocApp1.this.B), MdocApp1.this.A, MdocApp1.this.B *: // EmptyTuple.type] // -// val gridWrongValue = grid.set(Axis[A].at(0))(grid) -// ^^^^ +// val gridWrongValue = grid.updated(Axis[A].at(0))(grid) +// ^^^^ ``` ```scala // ERROR at runtime: a window with an Int start must lie inside the axis -grid.slice(Axis[A].at(3, 2)) +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:719) -// at dimwit.tensor.tensorops.StructuralExtensions$.slice(StructuralOps.scala:605) -// at dimwit.package$.slice(package.scala:80) -// at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:719) +// at dimwit.tensor.tensorops.StructuralExtensions$.windowExtents(StructuralOps.scala:742) +// at dimwit.tensor.tensorops.StructuralExtensions$.select(StructuralOps.scala:605) +// at dimwit.package$.select(package.scala:80) +// at repl.MdocSession$MdocApp17.$init$$$anonfun$3(AGENTS.md:723) ``` --- diff --git a/core/src/main/scala/dimwit/random/Random.scala b/core/src/main/scala/dimwit/random/Random.scala index 12d9af93..18b2c0d5 100644 --- a/core/src/main/scala/dimwit/random/Random.scala +++ b/core/src/main/scala/dimwit/random/Random.scala @@ -73,7 +73,7 @@ object Random: * Array(3, 4), * Array(5, 6)) ) * val perm = Random.permutation(Axis[Row] -> 3)(key) - * val shuffled = data.slice(Axis[Row].at(perm)) + * val shuffled = data(Axis[Row].at(perm)) * }}} * * @param dim diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala index 990eca31..eaa96f2a 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala @@ -320,7 +320,7 @@ private[dimwit] object StructuralOps: val headTensor = Tensor[NewShape, V](currentArr)(using newLabelsWitness) headTensor *: tailMaker(arrays.tail, compLabels.tail, originalLabels, splitIndex) -/** Extension methods for structural operations, e.g. `t.transpose`, `t.slice(...)` or `t.rearrange(...)`. */ +/** Extension methods for structural operations, e.g. `t.transpose`, `t.select(...)` or `t.rearrange(...)`. */ private[dimwit] object StructuralExtensions: import StructuralOps.* @@ -530,136 +530,159 @@ private[dimwit] object StructuralExtensions: val res = Jax.jnp.split(tensor.jaxValue, chunkSize, axis = axisIndex.index).as[Seq[Jax.PyDynamic]] res.map(x => Tensor[T, V](x)) - /** Slice the tensor at an index along an axis, removing the axis. + /** Select the tensor at an index along an axis, removing the axis. * * @param selector of the form `Axis[L].at(index)` */ - def slice[L](selector: AxisAtIndex[L])(using + def select[L](selector: AxisAtIndex[L])(using ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes] ): Tensor[ev.RemainingAxes, V] = Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(ev.index, py.Any.from(selector.index)))) - /** Slice the tensor at an index given by a tensor along an axis, removing the axis. + /** Select the tensor at an index given by a tensor along an axis, removing the axis. * The index may be traced, e.g. a loop index under `jit`. * * @param selector of the form `Axis[L].at(index: Tensor0[Int32])` */ - def slice[L](selector: AxisAtTensorIndex[L])(using + def select[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes] ): Tensor[ev.RemainingAxes, V] = Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(ev.index, selector.index.jaxValue))) - /** Slice the tensor at a range along an axis, keeping the axis with the length of the range as extent. + /** Select the tensor at a range along an axis, keeping the axis with the length of the range as extent. * * @param selector of the form `Axis[L].at(range)` */ - def slice[L](selector: AxisAtRange[L])(using + def select[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T] ): Tensor[T, V] = Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, tensor.pyRange(selector.range, axisIndex.index)))) - /** Slice the tensor at a sequence of indices along an axis, keeping the axis with the number of indices as extent. + /** Select the tensor at a sequence of indices along an axis, keeping the axis with the number of indices as extent. * * @param selector of the form `Axis[L].at(Seq(i, j, ...))` */ - def slice[L](selector: AxisAtIndices[L])(using + def select[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T] ): Tensor[T, V] = Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices)))) - /** Slice the tensor at a tuple of indices along an axis, keeping the axis with the number of indices as extent. + /** Select the tensor at a tuple of indices along an axis, keeping the axis with the number of indices as extent. * * @param selector of the form `Axis[L].at((i, j, ...))` */ - def slice[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using + def select[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T] ): Tensor[T, V] = Tensor(tensor.jaxValue.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(selector.indices.toList.asInstanceOf[List[Int]])))) - /** Slice the tensor at the indices of a vector along an axis, like `jnp.take`. + /** Select the tensor at the indices of a vector along an axis, like `jnp.take`. * The axis is replaced by the axis of the indices. Out-of-bounds indices yield NaN for floating tensors. * * @param selector of the form `Axis[L].at(indices: Tensor1[L2, Int32])` */ - def slice[L, L2](selector: AxisAtTensorIndices[L, L2])(using + def select[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[T, L, L2], labels: Labels[ev.NewShape] ): Tensor[ev.NewShape, V] = Tensor(Jax.jnp.take(tensor.jaxValue, selector.indices.jaxValue, axis = ev.index)) - /** Slice a window of static size along an axis, like `jax.lax.dynamic_slice`. + /** Select a window of static size along an axis, like `jax.lax.dynamic_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 window, which keeps the axis with extent `windowSize`. */ - def slice[L](selector: AxisAtWindow[L])(using + def select[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T] ): Tensor[T, V] = val window = tensor.windowExtents(selector, axisIndex.index) Tensor(Jax.lax.dynamic_slice(tensor.jaxValue, tensor.windowStarts(selector, axisIndex.index), window.toPythonCopy)) - /** Overwrite the tensor at an index along an axis. + // `tensor(selector)` is `tensor.select(selector)`, as `seq(i)` reads a Scala collection. + + def apply[L](selector: AxisAtIndex[L])(using ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes]): Tensor[ev.RemainingAxes, V] = + tensor.select(selector)(using ev, labels) + + def apply[L](selector: AxisAtTensorIndex[L])(using ev: AxisRemover[T, L], labels: Labels[ev.RemainingAxes]): Tensor[ev.RemainingAxes, V] = + tensor.select(selector)(using ev, labels) + + def apply[L](selector: AxisAtRange[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): Tensor[T, V] = + tensor.select(selector) + + def apply[L](selector: AxisAtIndices[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): Tensor[T, V] = + tensor.select(selector) + + def apply[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T]): Tensor[T, V] = + tensor.select(selector) + + def apply[L, L2](selector: AxisAtTensorIndices[L, L2])(using ev: AxisReplacer[T, L, L2], labels: Labels[ev.NewShape]): Tensor[ev.NewShape, V] = + tensor.select(selector)(using ev, labels) + + 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. * * @param selector of the form `Axis[L].at(index)` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L](selector: AxisAtIndex[L])(using + 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)) - /** Overwrite 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])` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L](selector: AxisAtTensorIndex[L])(using + 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)) - /** Overwrite 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)` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L](selector: AxisAtRange[L])(using + 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)) - /** Overwrite 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, ...))` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L](selector: AxisAtIndices[L])(using + 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)) - /** Overwrite 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, ...))` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using + def updated[L, I <: NonEmptyTuple](selector: AxisAtTupleIndices[L, I])(using axisIndex: AxisIndex[T, L], labels: Labels[T] )(value: Tensor[T, V]): Tensor[T, V] = @@ -667,26 +690,26 @@ private[dimwit] object StructuralExtensions: requireSelectionShape(value, tensor.shape.dimensions.updated(axisIndex.index, indices.length)) Tensor(tensor.jaxValue.at.itemAt(tensor.indexAlong(axisIndex.index, pyIndices(indices))).set(value.jaxValue)) - /** Overwrite 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])` - * @param value the new values, with the shape of `tensor.slice(selector)` + * @param value the new values, with the shape of `tensor.select(selector)` */ - def set[L, L2](selector: AxisAtTensorIndices[L, L2])(using + 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)) - /** Overwrite 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)` - * @param value the new window, with the shape of `tensor.slice(selector)` + * @param value the new window, with the shape of `tensor.select(selector)` */ - def set[L](selector: AxisAtWindow[L])(using + def updated[L](selector: AxisAtWindow[L])(using axisIndex: AxisIndex[T, L], labels: Labels[T] )(value: Tensor[T, V]): Tensor[T, V] = diff --git a/core/src/main/scala/dimwit/tensor/tensorops/Tensor0Extensions.scala b/core/src/main/scala/dimwit/tensor/tensorops/Tensor0Extensions.scala index ec743505..6a691c0a 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/Tensor0Extensions.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/Tensor0Extensions.scala @@ -11,7 +11,7 @@ private[dimwit] object Tensor0Extensions: """ | Cannot convert a JAX Tracer to a scalar value. Tensor0 is part of a JAX computation graph (e.g., inside vmap or a jitted function). | Common mistakes leading to this error: - | - calling .slice(t0.item) rather than .slice(t0); breaking the computation graph unintentionally. + | - calling t(Axis[L].at(t0.item)) rather than t(Axis[L].at(t0)); breaking the computation graph unintentionally. |""".stripMargin ) diff --git a/core/src/test/scala/dimwit/linalg/LinearAlgebraSuite.scala b/core/src/test/scala/dimwit/linalg/LinearAlgebraSuite.scala index e7c656c5..224002c9 100644 --- a/core/src/test/scala/dimwit/linalg/LinearAlgebraSuite.scala +++ b/core/src/test/scala/dimwit/linalg/LinearAlgebraSuite.scala @@ -124,7 +124,7 @@ class LinearAlgebraSuite extends DimwitTest: it("R is upper triangular: lower-left element is zero"): val (_, r) = LinearAlgebra.qr(qrMat, Axis[LBasis]) - r.slice(Axis[LBasis].at(1)).slice(Axis[Prime[A]].at(0)).item shouldBe 0.0f +- 1e-5f + r.select(Axis[LBasis].at(1)).select(Axis[Prime[A]].at(0)).item shouldBe 0.0f +- 1e-5f it("Frobenius norm is preserved: ||A||_F = ||R||_F (since Q is orthogonal)"): val (_, r) = LinearAlgebra.qr(qrMat, Axis[LBasis]) diff --git a/core/src/test/scala/dimwit/random/RandomSuite.scala b/core/src/test/scala/dimwit/random/RandomSuite.scala index 113fa724..13c92431 100644 --- a/core/src/test/scala/dimwit/random/RandomSuite.scala +++ b/core/src/test/scala/dimwit/random/RandomSuite.scala @@ -35,7 +35,7 @@ class RandomSuite extends DimwitTest: val splitKeys = key.split(n) for i <- 0 until n do - val tensorKey = tensorKeys.slice(Axis[Samples].at(i)).item + val tensorKey = tensorKeys.select(Axis[Samples].at(i)).item val splitKey = splitKeys(i) tensorKey should equal(splitKey) @@ -99,7 +99,7 @@ class RandomSuite extends DimwitTest: )) val rowPerm = Random.permutation(Axis[Row] -> 4)(key) - val shuffled = original.slice(Axis[Row].at(rowPerm)) + val shuffled = original.select(Axis[Row].at(rowPerm)) shuffled.shape should equal(original.shape) shuffled.sum.item shouldBe original.sum.item // sum should be unchanged @@ -107,7 +107,7 @@ class RandomSuite extends DimwitTest: // Check that each row in shuffled exists in original by comparing row sums // Original row sums are: [3, 12, 21, 30] val shuffledRowSums = (0 until 4).map { i => - shuffled.slice(Axis[Row].at(i)).sum.item + shuffled.select(Axis[Row].at(i)).sum.item } val expectedRowSums = Set(3, 12, 21, 30) shuffledRowSums.toSet shouldBe expectedRowSums diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala index 3d3f843b..68d6173f 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -164,9 +164,9 @@ class TensorCreationSuite extends DimwitTest: typeCheckErrors("Tensor1(Axis[A]).fromRange(0 until 3, VType[Float32])") should not be empty typeCheckErrors("Tensor1(Axis[A], VType[Float32]).fromRange(0 until 3)") should not be empty - it("can be consumed as gather indices by slice"): + it("can be consumed as gather indices by select"): val t = Tensor1(Axis[A]).fromArray(Array(10.0f, 20.0f, 30.0f)) - t.slice(Axis[A].at(Tensor1(Axis[B]).fromRange(0 until 3))) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) + t.select(Axis[A].at(Tensor1(Axis[B]).fromRange(0 until 3))) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) describe("linspace"): diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala index ed82688f..766e5a8d 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala @@ -162,10 +162,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = AB.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice(Axis[C].at(0)).slice(Axis[D].at(0)) should approxEqual(AB) - res.slice(Axis[C].at(1)).slice(Axis[D].at(0)) should approxEqual(AB) - res.slice(Axis[C].at(0)).slice(Axis[D].at(1)) should approxEqual(AB) - res.slice(Axis[C].at(1)).slice(Axis[D].at(1)) should approxEqual(AB) + res.select(Axis[C].at(0)).select(Axis[D].at(0)) should approxEqual(AB) + res.select(Axis[C].at(1)).select(Axis[D].at(0)) should approxEqual(AB) + res.select(Axis[C].at(0)).select(Axis[D].at(1)) should approxEqual(AB) + res.select(Axis[C].at(1)).select(Axis[D].at(1)) should approxEqual(AB) it("BC broadcastTo ABCD"): val BC = Tensor(Shape(Axis[B] -> 2, Axis[C] -> 2)).fromArray( @@ -173,10 +173,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = BC.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice(Axis[A].at(0)).slice(Axis[D].at(0)) should approxEqual(BC) - res.slice(Axis[A].at(1)).slice(Axis[D].at(0)) should approxEqual(BC) - res.slice(Axis[A].at(0)).slice(Axis[D].at(1)) should approxEqual(BC) - res.slice(Axis[A].at(1)).slice(Axis[D].at(1)) should approxEqual(BC) + res.select(Axis[A].at(0)).select(Axis[D].at(0)) should approxEqual(BC) + res.select(Axis[A].at(1)).select(Axis[D].at(0)) should approxEqual(BC) + res.select(Axis[A].at(0)).select(Axis[D].at(1)) should approxEqual(BC) + res.select(Axis[A].at(1)).select(Axis[D].at(1)) should approxEqual(BC) it("CD broadcastTo ABCD"): val CD = Tensor(Shape(Axis[C] -> 2, Axis[D] -> 2)).fromArray( @@ -184,10 +184,10 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) val res = CD.broadcastTo(tABCD.shape) res.shape shouldEqual tABCD.shape - res.slice(Axis[A].at(0)).slice(Axis[B].at(0)) should approxEqual(CD) - res.slice(Axis[A].at(1)).slice(Axis[B].at(0)) should approxEqual(CD) - res.slice(Axis[A].at(0)).slice(Axis[B].at(1)) should approxEqual(CD) - res.slice(Axis[A].at(1)).slice(Axis[B].at(1)) should approxEqual(CD) + res.select(Axis[A].at(0)).select(Axis[B].at(0)) should approxEqual(CD) + res.select(Axis[A].at(1)).select(Axis[B].at(0)) should approxEqual(CD) + res.select(Axis[A].at(0)).select(Axis[B].at(1)) should approxEqual(CD) + res.select(Axis[A].at(1)).select(Axis[B].at(1)) should approxEqual(CD) describe("Disallow"): diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 0485dbcf..190bd437 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -435,9 +435,9 @@ class TensorOpsStructureSuite extends DimwitTest: val joined = concatenate(Seq(part1, part2, part3), Axis[B]) joined.axes shouldBe List("A", "B") joined.shape(Axis[B]) shouldBe (part1.shape(Axis[B]) + part2.shape(Axis[B]) + part3.shape(Axis[B])) - joined.slice(Axis[B].at(0 until part1.shape(Axis[B]))) should approxEqual(part1) - joined.slice(Axis[B].at(part1.shape(Axis[B]) until (part1.shape(Axis[B]) + part2.shape(Axis[B])))) should approxEqual(part2) - joined.slice(Axis[B].at((part1.shape(Axis[B]) + part2.shape(Axis[B])) until (part1.shape(Axis[B]) + part2.shape(Axis[B]) + part3.shape(Axis[B])))) should approxEqual(part3) + joined.select(Axis[B].at(0 until part1.shape(Axis[B]))) should approxEqual(part1) + joined.select(Axis[B].at(part1.shape(Axis[B]) until (part1.shape(Axis[B]) + part2.shape(Axis[B])))) should approxEqual(part2) + joined.select(Axis[B].at((part1.shape(Axis[B]) + part2.shape(Axis[B])) until (part1.shape(Axis[B]) + part2.shape(Axis[B]) + part3.shape(Axis[B])))) should approxEqual(part3) it("concatenate2 different axes"): val part1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) @@ -445,8 +445,8 @@ class TensorOpsStructureSuite extends DimwitTest: val joined = concatenate(part1, part2) joined.axes shouldBe List("A", "B+C") joined.shape(Axis[B |+| C]) shouldBe (part1.shape(Axis[B]) + part2.shape(Axis[C])) - joined.slice(Axis[B |+| C].at(0 until part1.shape(Axis[B]))) should approxEqual(part1.relabel(Axis[B].as(Axis[B |+| C]))) - joined.slice(Axis[B |+| C].at(part1.shape(Axis[B]) until (part1.shape(Axis[B]) + part2.shape(Axis[C])))) should approxEqual(part2.relabel(Axis[C].as(Axis[B |+| C]))) + joined.select(Axis[B |+| C].at(0 until part1.shape(Axis[B]))) should approxEqual(part1.relabel(Axis[B].as(Axis[B |+| C]))) + joined.select(Axis[B |+| C].at(part1.shape(Axis[B]) until (part1.shape(Axis[B]) + part2.shape(Axis[C])))) should approxEqual(part2.relabel(Axis[C].as(Axis[B |+| C]))) describe("Deconcatenation"): @@ -493,89 +493,101 @@ class TensorOpsStructureSuite extends DimwitTest: describe("selectors"): it("are applied one at a time, not as a tuple"): - "t3.slice((Axis[A].at(0), Axis[B].at(0)))" shouldNot compile - "t3.set((Axis[A].at(0), Axis[B].at(0)))(Tensor1(Axis[C]).fromArray(Array(0f)))" shouldNot compile + "t3.select((Axis[A].at(0), Axis[B].at(0)))" shouldNot compile + "t3.updated((Axis[A].at(0), Axis[B].at(0)))(Tensor1(Axis[C]).fromArray(Array(0f)))" shouldNot compile + + it("can be applied to the tensor like select"): + val ab = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(3, 4)((a, b) => 10f * a + b)) + val indices = Tensor1(Axis[C]).fromArray(Array(2, 0)) + ab(Axis[A].at(1)) should approxEqual(ab.select(Axis[A].at(1))) + ab(Axis[A].at(Tensor0(1))) should approxEqual(ab.select(Axis[A].at(Tensor0(1)))) + ab(Axis[A].at(0 until 2)) should approxEqual(ab.select(Axis[A].at(0 until 2))) + ab(Axis[A].at(Seq(2, 0))) should approxEqual(ab.select(Axis[A].at(Seq(2, 0)))) + ab(Axis[A].at((2, 0))) should approxEqual(ab.select(Axis[A].at((2, 0)))) + ab(Axis[B].at(indices)) should approxEqual(ab.select(Axis[B].at(indices))) + ab(Axis[B].at(Tensor0(1), 2)) should approxEqual(ab.select(Axis[B].at(Tensor0(1), 2))) + ab(Axis[A].at(1))(Axis[B].at(2)).item shouldBe 12f it("reject a value whose shape differs from the selection"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) - the[IllegalArgumentException] thrownBy t.set(Axis[A].at(0 to 1))(t) should have message + the[IllegalArgumentException] thrownBy t.updated(Axis[A].at(0 to 1))(t) should have message "requirement failed: Cannot set a value of shape Shape(A -> 3) into a selection of shape (A -> 2)" - describe("set function"): + describe("updated function"): describe("AxisAtIndex"): - it("set single value (scalar index)"): + it("updates a single value (scalar index)"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - t = t.set(Axis[A].at(0))(t.slice(Axis[A].at(0)).set(Axis[B].at(0))(5f)) + t = t.updated(Axis[A].at(0))(t.select(Axis[A].at(0)).updated(Axis[B].at(0))(5f)) t should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(5.0f, 2.0f), Array(3.0f, 4.0f)))) - it("set single value (tensor index)"): + it("updates a single value (tensor index)"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - t = t.set(Axis[A].at(Tensor0(1)))(t.slice(Axis[A].at(Tensor0(1))).set(Axis[B].at(Tensor0(1)))(5f)) + t = t.updated(Axis[A].at(Tensor0(1)))(t.select(Axis[A].at(Tensor0(1))).updated(Axis[B].at(Tensor0(1)))(5f)) t should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 5.0f)))) - it("set a sub-vector"): + it("updates a sub-vector"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - t = t.set(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(5.0f, 6.0f))) + t = t.updated(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(5.0f, 6.0f))) t should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(5.0f, 6.0f), Array(3.0f, 4.0f)))) it("auto-broadcasting not supported"): var t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - "t = t.set(Axis[A].at(0))(Tensor0(3.0f))" shouldNot compile + "t = t.updated(Axis[A].at(0))(Tensor0(3.0f))" shouldNot compile describe("AxisAtRange"): - it("set at a range"): + it("updates at a range"): var t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) - val t2 = t.set(Axis[A].at(0 to 1))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) + val t2 = t.updated(Axis[A].at(0 to 1))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) t2 should approxEqual(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f, 3.0f))) - val t3 = t.set(Axis[A].at(1 to 2))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) + val t3 = t.updated(Axis[A].at(1 to 2))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) t3 should approxEqual(Tensor1(Axis[A]).fromArray(Array(1.0f, 7.0f, 8.0f))) - it("slice at a range"): + it("select at a range"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) - t.slice(Axis[A].at(1 to 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) - t.slice(Axis[A].at(1 until 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) - t.slice(Axis[A].at(0 until 4 by 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(1.0f, 3.0f))) + t.select(Axis[A].at(1 to 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) + t.select(Axis[A].at(1 until 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f))) + t.select(Axis[A].at(0 until 4 by 2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(1.0f, 3.0f))) - it("slice at a descending range"): + it("select at a descending range"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) - t.slice(Axis[A].at(3 to 1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) - t.slice(Axis[A].at(3 until 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) + t.select(Axis[A].at(3 to 1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) + t.select(Axis[A].at(3 until 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f))) // both spellings of a descending range down to index 0 - t.slice(Axis[A].at(3 to 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) - t.slice(Axis[A].at(3 until -1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) - t.slice(Axis[A].at(3 to 0 by -2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 2.0f))) + t.select(Axis[A].at(3 to 0 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) + t.select(Axis[A].at(3 until -1 by -1)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 3.0f, 2.0f, 1.0f))) + t.select(Axis[A].at(3 to 0 by -2)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 2.0f))) - it("slice rejects a range that leaves the axis"): + it("select rejects a range that leaves the axis"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) - the[IllegalArgumentException] thrownBy t.slice(Axis[A].at(3 until -10000 by -1)) should have message + the[IllegalArgumentException] thrownBy t.select(Axis[A].at(3 until -10000 by -1)) should have message "requirement failed: Range 3 until -10000 by -1 is out of bounds for axis of size 4" - an[IllegalArgumentException] should be thrownBy t.slice(Axis[A].at(3 until -3 by -1)) - an[IllegalArgumentException] should be thrownBy t.slice(Axis[A].at(0 until 10)) + an[IllegalArgumentException] should be thrownBy t.select(Axis[A].at(3 until -3 by -1)) + an[IllegalArgumentException] should be thrownBy t.select(Axis[A].at(0 until 10)) // a very negative end is fine as long as the step steps over the invalid indices - t.slice(Axis[A].at(3 until -3 by -3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 1.0f))) + t.select(Axis[A].at(3 until -3 by -3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(4.0f, 1.0f))) - it("slice at an empty range yields a zero-length axis"): + it("select at an empty range yields a zero-length axis"): val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) - val empty = t.slice(Axis[A].at(0 until 0)) + val empty = t.select(Axis[A].at(0 until 0)) empty.axes shouldBe List("A", "B") empty.shape(Axis[A]) shouldBe 0 empty.shape(Axis[B]) shouldBe 2 - t.slice(Axis[A].at(2 until 2)).shape(Axis[A]) shouldBe 0 - t.slice(Axis[A].at(1 until 1)).shape(Axis[A]) shouldBe 0 + t.select(Axis[A].at(2 until 2)).shape(Axis[A]) shouldBe 0 + t.select(Axis[A].at(1 until 1)).shape(Axis[A]) shouldBe 0 - it("set at an empty range is a no-op"): + it("updates at an empty range is a no-op"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) val empty = Tensor1(Axis[A]).fromArray(Array.empty[Float]) - t.set(Axis[A].at(1 until 1))(empty) should approxEqual(t) + t.updated(Axis[A].at(1 until 1))(empty) should approxEqual(t) describe("AxisAtIndices"): - it("set at a seq"): + it("updates at a seq"): var t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) - val t2 = t.set(Axis[A].at(List(0, 2)))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) + val t2 = t.updated(Axis[A].at(List(0, 2)))(Tensor1(Axis[A]).fromArray(Array(7.0f, 8.0f))) t2 should approxEqual(Tensor1(Axis[A]).fromArray(Array(7.0f, 2.0f, 8.0f))) describe("AxisAtTensorIndices"): @@ -584,24 +596,24 @@ class TensorOpsStructureSuite extends DimwitTest: val ab = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(2, 3)((a, b) => 10f * a + b)) val indices = Tensor1(Axis[C]).fromArray(Array(2, 0, 2, 1)) - it("slices at the indices and replaces the axis by the axis of the indices"): - val sliced = ab.slice(Axis[B].at(indices)) + it("selects at the indices and replaces the axis by the axis of the indices"): + val sliced = ab.select(Axis[B].at(indices)) sliced.axes shouldBe List("A", "C") sliced should approxEqual(Tensor2(Axis[A], Axis[C]).fromArray(Array(Array(2f, 0f, 2f, 1f), Array(12f, 10f, 12f, 11f)))) it("keeps the axis when the indices have the same axis"): val permutation = Tensor1(Axis[B]).fromArray(Array(2, 0, 1)) - ab.slice(Axis[B].at(permutation)) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(12f, 10f, 11f)))) + ab.select(Axis[B].at(permutation)) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(12f, 10f, 11f)))) - it("sets at the indices"): + it("updates at the indices"): val zeros = Tensor(ab.shape).fill(0f) val positions = Tensor1(Axis[C]).fromArray(Array(2, 0)) val values = Tensor2(Axis[A], Axis[C]).fromArray(Array(Array(1f, 2f), Array(3f, 4f))) - zeros.set(Axis[B].at(positions))(values) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(4f, 0f, 3f)))) + zeros.updated(Axis[B].at(positions))(values) should approxEqual(Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(2f, 0f, 1f), Array(4f, 0f, 3f)))) - it("set is the inverse of slice for distinct indices"): + it("updated is the inverse of select for distinct indices"): val permutation = Tensor1(Axis[B]).fromArray(Array(2, 0, 1)) - Tensor(ab.shape).fill(0f).set(Axis[B].at(permutation))(ab.slice(Axis[B].at(permutation))) should approxEqual(ab) + Tensor(ab.shape).fill(0f).updated(Axis[B].at(permutation))(ab.select(Axis[B].at(permutation))) should approxEqual(ab) describe("roll"): it("tensor1"): @@ -622,88 +634,88 @@ class TensorOpsStructureSuite extends DimwitTest: Array.tabulate(2, 5, 3)((a, b, c) => 100f * a + 10f * b + c) ) - describe("slice"): + describe("select"): - it("slices a window of a vector"): + it("selects a window of a vector"): val t = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f)) - t.slice(Axis[A].at(Tensor0(2), 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f, 5.0f))) + t.select(Axis[A].at(Tensor0(2), 3)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f, 5.0f))) - it("slices a window along a middle axis and keeps the axis"): - val window: Tensor3[A, B, C, Float32] = abc.slice(Axis[B].at(Tensor0(1), 2)) + it("selects a window along a middle axis and keeps the axis"): + val window: Tensor3[A, B, C, Float32] = abc.select(Axis[B].at(Tensor0(1), 2)) window.axes shouldBe List("A", "B", "C") window.shape.dimensions shouldBe List(2, 2, 3) - window should approxEqual(abc.slice(Axis[B].at(1 until 3))) + window should approxEqual(abc.select(Axis[B].at(1 until 3))) it("clamps the start so that the window fits"): - abc.slice(Axis[B].at(Tensor0(4), 3)) should approxEqual(abc.slice(Axis[B].at(2 until 5))) + abc.select(Axis[B].at(Tensor0(4), 3)) should approxEqual(abc.select(Axis[B].at(2 until 5))) it("chains with other selectors"): - val window = abc.slice(Axis[B].at(Tensor0(1), 2)).slice(Axis[C].at(0)) + val window = abc.select(Axis[B].at(Tensor0(1), 2)).select(Axis[C].at(0)) window.axes shouldBe List("A", "B") - window should approxEqual(abc.slice(Axis[B].at(1 until 3)).slice(Axis[C].at(0))) + window should approxEqual(abc.select(Axis[B].at(1 until 3)).select(Axis[C].at(0))) it("rejects an axis that the tensor does not have"): - "abc.slice(Axis[D].at(Tensor0(1), 2))" shouldNot compile + "abc.select(Axis[D].at(Tensor0(1), 2))" shouldNot compile it("works with a traced start under jit"): - val sliceAt = jit((start: Tensor0[Int32]) => abc.slice(Axis[B].at(start, 2))) - sliceAt(Tensor0(0)) should approxEqual(abc.slice(Axis[B].at(0 until 2))) - sliceAt(Tensor0(3)) should approxEqual(abc.slice(Axis[B].at(3 until 5))) + val sliceAt = jit((start: Tensor0[Int32]) => abc.select(Axis[B].at(start, 2))) + sliceAt(Tensor0(0)) should approxEqual(abc.select(Axis[B].at(0 until 2))) + sliceAt(Tensor0(3)) should approxEqual(abc.select(Axis[B].at(3 until 5))) it("is differentiable"): - val grad = Autodiff.grad((t: Tensor3[A, B, C, Float32]) => t.slice(Axis[B].at(Tensor0(1), 2)).sum) + val grad = Autodiff.grad((t: Tensor3[A, B, C, Float32]) => t.select(Axis[B].at(Tensor0(1), 2)).sum) val expected = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( Array.tabulate(2, 5, 3)((_, b, _) => if b == 1 || b == 2 then 1f else 0f) ) grad(abc).value should approxEqual(expected) - it("slices a window with an Int start"): - abc.slice(Axis[B].at(1, 2)) should approxEqual(abc.slice(Axis[B].at(1 until 3))) + it("selects a window with an Int start"): + abc.select(Axis[B].at(1, 2)) should approxEqual(abc.select(Axis[B].at(1 until 3))) it("rejects a window with an Int start that does not fit"): - the[IllegalArgumentException] thrownBy abc.slice(Axis[B].at(4, 3)) should have message + the[IllegalArgumentException] thrownBy abc.select(Axis[B].at(4, 3)) should have message "requirement failed: Window 4 until 7 is out of bounds for axis B of size 5" it("rejects a window larger than the axis"): - the[IllegalArgumentException] thrownBy abc.slice(Axis[B].at(Tensor0(0), 6)) should have message + the[IllegalArgumentException] thrownBy abc.select(Axis[B].at(Tensor0(0), 6)) should have message "requirement failed: Window of size 6 does not fit axis B of size 5" - describe("set"): + describe("updated"): val zeros = Tensor(abc.shape).fill(0f) val block = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( Array.tabulate(2, 2, 3)((a, b, c) => 1f + 100f * a + 10f * b + c) ) - it("sets a window along a middle axis"): - zeros.set(Axis[B].at(Tensor0(1), 2))(block) should approxEqual(zeros.set(Axis[B].at(1 until 3))(block)) + it("updates a window along a middle axis"): + zeros.updated(Axis[B].at(Tensor0(1), 2))(block) should approxEqual(zeros.updated(Axis[B].at(1 until 3))(block)) it("clamps the start so that the window fits"): - zeros.set(Axis[B].at(Tensor0(4), 2))(block) should approxEqual(zeros.set(Axis[B].at(3 until 5))(block)) + zeros.updated(Axis[B].at(Tensor0(4), 2))(block) should approxEqual(zeros.updated(Axis[B].at(3 until 5))(block)) it("works with a traced start under jit"): - val setAt = jit((start: Tensor0[Int32]) => zeros.set(Axis[B].at(start, 2))(block)) - setAt(Tensor0(0)) should approxEqual(zeros.set(Axis[B].at(0 until 2))(block)) - setAt(Tensor0(3)) should approxEqual(zeros.set(Axis[B].at(3 until 5))(block)) + val setAt = jit((start: Tensor0[Int32]) => zeros.updated(Axis[B].at(start, 2))(block)) + setAt(Tensor0(0)) should approxEqual(zeros.updated(Axis[B].at(0 until 2))(block)) + setAt(Tensor0(3)) should approxEqual(zeros.updated(Axis[B].at(3 until 5))(block)) - it("sets a window with an Int start"): - zeros.set(Axis[B].at(1, 2))(block) should approxEqual(zeros.set(Axis[B].at(1 until 3))(block)) + it("updates a window with an Int start"): + zeros.updated(Axis[B].at(1, 2))(block) should approxEqual(zeros.updated(Axis[B].at(1 until 3))(block)) it("fills a buffer row by row"): val filled = (0 until 5).foldLeft(zeros): (buffer, b) => - buffer.set(Axis[B].at(Tensor0(b), 1))(abc.slice(Axis[B].at(b until b + 1))) + buffer.updated(Axis[B].at(Tensor0(b), 1))(abc.select(Axis[B].at(b until b + 1))) filled should approxEqual(abc) it("is differentiable with respect to the value"): - val grad = Autodiff.grad((value: Tensor3[A, B, C, Float32]) => (zeros.set(Axis[B].at(Tensor0(1), 2))(value) * abc).sum) - grad(block).value should approxEqual(abc.slice(Axis[B].at(1 until 3))) + val grad = Autodiff.grad((value: Tensor3[A, B, C, Float32]) => (zeros.updated(Axis[B].at(Tensor0(1), 2))(value) * abc).sum) + grad(block).value should approxEqual(abc.select(Axis[B].at(1 until 3))) it("rejects a value whose extent differs from the window size"): - the[IllegalArgumentException] thrownBy zeros.set(Axis[B].at(Tensor0(1), 3))(block) should have message + the[IllegalArgumentException] thrownBy zeros.updated(Axis[B].at(Tensor0(1), 3))(block) should have message "requirement failed: Cannot set a value of shape Shape(A -> 2, B -> 2, C -> 3) into a selection of shape (A -> 2, B -> 3, C -> 3)" it("does not broadcast a value without the window axis"): - "zeros.set(Axis[B].at(Tensor0(1), 2))(block.slice(Axis[B].at(0)))" shouldNot compile + "zeros.updated(Axis[B].at(Tensor0(1), 2))(block.select(Axis[B].at(0)))" shouldNot compile describe("chunk function"): diff --git a/docs/quickstart.md b/docs/quickstart.md index e5f378cf..97771dbb 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -300,7 +300,7 @@ To unflatten a partially flattened tensor, we need to specify the axis that we w ```scala partiallyFlattened.unflatten(Axis[B |*| C], Shape(Axis[B] -> 2, Axis[C] -> 4)) ``` -#### Concatenating, splitting and slicing tensors +#### Concatenating, splitting and selecting tensors Flatten and unflatten takes a single tensor and transform its shape. In contrast, concatenate and split take multiple tensors and combine them into a single tensor or split a single tensor into multiple tensors. Given two tensors with the same shape except for one axis, we can concatenate them along that axis as follows: @@ -321,18 +321,31 @@ If we want a split into several tensors, we can specify the split points as foll val (splt1, splt2, splt3) = concatenated.split(Axis[A].at((1, 2))) ``` -The method `slice` works in a similar way to split, but instead of splitting the tensor into several tensors, -it returns a single tensor that is a slice of the original tensor. For example, we can extract the slice -at index 1 along axis A as follows: +The method `select` works in a similar way to split, but instead of splitting the tensor into several tensors, +it returns a single tensor with the selected part of the original tensor. For example, we can select +index 1 along axis A as follows: ```scala - val sliced1 : Tensor1[B, Float32]= concatenated.slice(Axis[A].at(1)) + val selected1 : Tensor1[B, Float32]= concatenated.select(Axis[A].at(1)) ``` -As for split, we can also provide multiple a tuple (or sequence) of indices, which will then select all slices in the tuple. +As for Scala collections, we can also apply the tensor to the selector directly, which is the same as `select`: ```scala - val slicedMultiple : Tensor2[A, B, Float32] = concatenated.slice(Axis[A].at((0, 2))) + val applied1 : Tensor1[B, Float32]= concatenated(Axis[A].at(1)) +``` + +As for split, we can also provide a tuple (or sequence) of indices, which will then select all indices in the tuple. + +```scala + val selectedMultiple : Tensor2[A, B, Float32] = concatenated.select(Axis[A].at((0, 2))) +``` + +To change the values at a selection, `updated` returns a copy of the tensor with new values there, +as `updated` does for Scala's immutable collections: + +```scala + val updated1 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1))(Tensor(Shape(Axis[B] -> 2)).fill(0.0f)) ``` #### Squeezing, Expanding and transposing axes diff --git a/examples/src/main/scala/dimwit/basic/KMeans.scala b/examples/src/main/scala/dimwit/basic/KMeans.scala index 9719a2f0..6747cf14 100644 --- a/examples/src/main/scala/dimwit/basic/KMeans.scala +++ b/examples/src/main/scala/dimwit/basic/KMeans.scala @@ -114,7 +114,7 @@ trait CenterBasedClustering( key: Random.Key ): Tensor2[Cluster, Dim, Float32] = val perm = Random.permutation(points.shape.extent(Axis[Point]))(key) - points.slice(Axis[Point].at(perm.slice(Axis[Point].at(0 until nClusters)))).relabel(Axis[Point].as(Axis[Cluster])) + points(Axis[Point].at(perm(Axis[Point].at(0 until nClusters)))).relabel(Axis[Point].as(Axis[Cluster])) /** Example implementation of the K-Means clustering algorithm, showcasing the use * of Named Tensors for clear and type-safe code. @@ -166,10 +166,10 @@ object KMeans: println(" " + "─" * 65) (0 until k).foreach: c => - // Extract the c-th centroid row: slice along the named Cluster axis. - val center = centers.slice(Axis[Cluster].at(c)) - val cx = center.slice(Axis[Dim].at(0)).item - val cy = center.slice(Axis[Dim].at(1)).item + // Extract the c-th centroid row: select along the named Cluster axis. + val center = centers(Axis[Cluster].at(c)) + val cx = center(Axis[Dim].at(0)).item + val cy = center(Axis[Dim].at(1)).item // Cluster size: count points whose assignment equals c. // elementEquals produces [Point, Bool]; asFloat32 then sum gives a count. diff --git a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala index 471e60de..df4e2e7e 100644 --- a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala +++ b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala @@ -86,14 +86,14 @@ object LogisticRegression: val perm = Random.permutation(dataInitial.shape.extent(Axis[Sample]))(shuffleKey) val testSplitRatio = 0.4f val splitIndex = (numSamples * testSplitRatio).toInt - val trainPerm = perm.slice(Axis[Sample].at(0 until splitIndex)) - val testPerm = perm.slice(Axis[Sample].at(splitIndex until numSamples)) + val trainPerm = perm(Axis[Sample].at(0 until splitIndex)) + val testPerm = perm(Axis[Sample].at(splitIndex until numSamples)) // Use the permutations to get our training and validation data - val trainingDataUnnormalized = dataInitial.slice(Axis[Sample].at(trainPerm)) - val valDataUnnormalized = dataInitial.slice(Axis[Sample].at(testPerm)) - val trainLabels = labelsInitial.slice(Axis[Sample].at(trainPerm)) - val valLabels = labelsInitial.slice(Axis[Sample].at(testPerm)) + val trainingDataUnnormalized = dataInitial(Axis[Sample].at(trainPerm)) + val valDataUnnormalized = dataInitial(Axis[Sample].at(testPerm)) + val trainLabels = labelsInitial(Axis[Sample].at(trainPerm)) + val valLabels = labelsInitial(Axis[Sample].at(testPerm)) def calcMeanAndStd(t: Tensor2[Sample, Feature, Float32]): (Tensor1[Feature, Float32], Tensor1[Feature, Float32]) = val mean = t.vmap(Axis[Feature])(_.mean) diff --git a/examples/src/main/scala/dimwit/basic/SIRSimulation.scala b/examples/src/main/scala/dimwit/basic/SIRSimulation.scala index 53554d4d..43e6d0cc 100644 --- a/examples/src/main/scala/dimwit/basic/SIRSimulation.scala +++ b/examples/src/main/scala/dimwit/basic/SIRSimulation.scala @@ -129,10 +129,10 @@ object SIRSimulation: val infectedOverTime: Tensor1[Time, Float32] = stack(trajectory.map(_.I), Axis[Time]).sum(Axis[Group]) - println(s"I(0) = ${infectedOverTime.slice(Axis[Time].at(0))}") - println(s"I(mid) = ${infectedOverTime.slice(Axis[Time].at(nSteps / 2))}") - println(s"I(end) = ${infectedOverTime.slice(Axis[Time].at(nSteps))}") + println(s"I(0) = ${infectedOverTime(Axis[Time].at(0))}") + println(s"I(mid) = ${infectedOverTime(Axis[Time].at(nSteps / 2))}") + println(s"I(end) = ${infectedOverTime(Axis[Time].at(nSteps))}") val finalInfectedByGroup: Tensor1[Group, Float32] = - stack(trajectory.map(_.I), Axis[Time]).slice(Axis[Time].at(nSteps)) + stack(trajectory.map(_.I), Axis[Time]).select(Axis[Time].at(nSteps)) println(s"Final infected by group: $finalInfectedByGroup") diff --git a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala index dcfd9817..1b912c4a 100644 --- a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala +++ b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala @@ -248,8 +248,7 @@ object VariationalAutoencoderExample: val vae = VariationalAutoencoder(trainedParams) /* Reconstructing images */ - val reconstructed = testImages - .slice(Axis[TestSample].at(0 until 64)) + val reconstructed = testImages(Axis[TestSample].at(0 until 64)) .vmap(Axis[TestSample]): sample => val (mean, logVar) = vae.encoder(sample.flatten) val latent = reparametrize(mean, logVar, dataKey) // TODo Key management diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index f9b97afd..ee51ba84 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -467,14 +467,14 @@ println(s"Transposed shape: ${transposed.shape}") val vec1d = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f)) val vec2d = vec1d.appendAxis(Axis[B]) // Add new axis B -// Slice at an index (removes the axis) -val firstRow = original.slice(Axis[A].at(0)) -val secondRow = original.slice(Axis[A].at(1)) +// Select at an index (removes the axis) +val firstRow = original.select(Axis[A].at(0)) +val secondRow = original.select(Axis[A].at(1)) -// Slice at a range (keeps the axis) +// Select at a range (keeps the axis) val data3d = Tensor(Shape3(Axis[A] -> 5, Axis[B] -> 3, Axis[C] -> 4)).fill(1.0f) -val middleSlice = data3d.slice(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 -println(s"Middle slice shape: ${middleSlice.shape}") // Shape(A -> 3, B -> 3, C -> 4) +val middleSelection = data3d.select(Axis[A].at(1 until 4)) // Takes indices 1, 2, 3 +println(s"Middle selection shape: ${middleSelection.shape}") // Shape(A -> 3, B -> 3, C -> 4) // Concatenate val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) @@ -482,37 +482,41 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) val concatenated = concatenate(t1, t2, Axis[A]) ``` -### Slicing and Setting +### Selecting and Updating -Every selection along an axis is an `Axis[L].at(...)` selector. `slice` reads the selection, -`set` returns a copy with the selection overwritten. A selector selects along exactly one axis; -to select along several axes, chain the calls. +Every selection along an axis is an `Axis[L].at(...)` selector. `select` reads the selection, and applying +the tensor to the selector, `t(selector)`, is the same as `t.select(selector)`, as `seq(i)` reads a Scala collection. +`updated` returns a copy with new values at the selection, as `updated` does for Scala's immutable collections. +A selector selects along exactly one axis; to select along several axes, chain the calls. ```scala mdoc:silent // Shape: A=4, B=3, the value at (a, b) is 10a + b val grid = Tensor2(Axis[A], Axis[B]).fromArray(Array.tabulate(4, 3)((a, b) => 10.0f * a + b)) // Index: removes the axis -val gridRow: Tensor1[B, Float32] = grid.slice(Axis[A].at(1)) // [10, 11, 12] +val gridRow: Tensor1[B, Float32] = grid.select(Axis[A].at(1)) // [10, 11, 12] // Range, sequence or tuple of indices: keeps the axis, with the number of indices as extent -val gridRows = grid.slice(Axis[A].at(1 until 3)) // rows 1, 2 -val gridPicked = grid.slice(Axis[A].at(Seq(3, 0))) // rows 3, 0 -val gridTupled = grid.slice(Axis[A].at((3, 0))) // the same, note the double parentheses +val gridRows = grid.select(Axis[A].at(1 until 3)) // rows 1, 2 +val gridPicked = grid.select(Axis[A].at(Seq(3, 0))) // rows 3, 0 +val gridTupled = grid.select(Axis[A].at((3, 0))) // the same, note the double parentheses // Vector of indices, like jnp.take: the axis is replaced by the axis of the indices val order = Tensor1(Axis[C]).fromArray(Array(2, 0, 2)) -val gridGathered: Tensor2[A, C, Float32] = grid.slice(Axis[B].at(order)) +val gridGathered: Tensor2[A, C, Float32] = grid.select(Axis[B].at(order)) // Window of static size, like jax.lax.dynamic_slice: keeps the axis with the window size as extent -val gridWindow = grid.slice(Axis[A].at(1, 2)) // rows 1, 2 +val gridWindow = grid.select(Axis[A].at(1, 2)) // rows 1, 2 // Several axes: chain the calls -val gridEntry: Tensor0[Float32] = grid.slice(Axis[A].at(1)).slice(Axis[B].at(2)) // 12 +val gridEntry: Tensor0[Float32] = grid.select(Axis[A].at(1)).select(Axis[B].at(2)) // 12 -// set takes a value with the shape of the selection -val gridZeroRow = grid.set(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) -val gridZeroWindow = grid.set(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) +// Apply: the same as select +val gridEntryApplied: Tensor0[Float32] = grid(Axis[A].at(1))(Axis[B].at(2)) // 12 + +// updated takes a value with the shape of the selection +val gridZeroRow = grid.updated(Axis[A].at(0))(Tensor1(Axis[B]).fromArray(Array(0.0f, 0.0f, 0.0f))) +val gridZeroWindow = grid.updated(Axis[A].at(1, 2))(Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 3)).fill(0.0f)) ``` The index of `Axis[L].at(index)` and the start of a window may also be a `Tensor0[Int32]`, also a traced one, @@ -521,24 +525,24 @@ e.g. a loop position under `jit`. Its value is then only known at runtime, but t ```scala mdoc:silent // Write a row of ones into a buffer at a position known only at runtime val writeRow = jit((buffer: Tensor2[A, B, Float32], position: Tensor0[Int32]) => - buffer.set(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) + buffer.updated(Axis[A].at(position, 1))(Tensor(Shape2(Axis[A] -> 1, Axis[B] -> 3)).fill(1.0f)) ) val gridWritten = writeRow(grid, Tensor0(2)) // row 2 is now ones ``` ```scala mdoc:fail // ERROR: a selector selects along one axis, chain the calls instead -val gridTuple = grid.slice((Axis[A].at(1), Axis[B].at(2))) +val gridTuple = grid.select((Axis[A].at(1), Axis[B].at(2))) ``` ```scala mdoc:fail // ERROR: an index removes the axis, so the value must not have it -val gridWrongValue = grid.set(Axis[A].at(0))(grid) +val gridWrongValue = grid.updated(Axis[A].at(0))(grid) ``` ```scala mdoc:crash // ERROR at runtime: a window with an Int start must lie inside the axis -grid.slice(Axis[A].at(3, 2)) +grid.select(Axis[A].at(3, 2)) ``` --- diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md index af10fe65..e0e4922d 100644 --- a/mdocs/docs/quickstart.md +++ b/mdocs/docs/quickstart.md @@ -273,7 +273,7 @@ To unflatten a partially flattened tensor, we need to specify the axis that we w ```scala mdoc:silent partiallyFlattened.unflatten(Axis[B |*| C], Shape(Axis[B] -> 2, Axis[C] -> 4)) ``` -#### Concatenating, splitting and slicing tensors +#### Concatenating, splitting and selecting tensors Flatten and unflatten takes a single tensor and transform its shape. In contrast, concatenate and split take multiple tensors and combine them into a single tensor or split a single tensor into multiple tensors. Given two tensors with the same shape except for one axis, we can concatenate them along that axis as follows: @@ -294,18 +294,31 @@ If we want a split into several tensors, we can specify the split points as foll val (splt1, splt2, splt3) = concatenated.split(Axis[A].at((1, 2))) ``` -The method `slice` works in a similar way to split, but instead of splitting the tensor into several tensors, -it returns a single tensor that is a slice of the original tensor. For example, we can extract the slice -at index 1 along axis A as follows: +The method `select` works in a similar way to split, but instead of splitting the tensor into several tensors, +it returns a single tensor with the selected part of the original tensor. For example, we can select +index 1 along axis A as follows: ```scala mdoc:silent - val sliced1 : Tensor1[B, Float32]= concatenated.slice(Axis[A].at(1)) + val selected1 : Tensor1[B, Float32]= concatenated.select(Axis[A].at(1)) ``` -As for split, we can also provide multiple a tuple (or sequence) of indices, which will then select all slices in the tuple. +As for Scala collections, we can also apply the tensor to the selector directly, which is the same as `select`: ```scala mdoc:silent - val slicedMultiple : Tensor2[A, B, Float32] = concatenated.slice(Axis[A].at((0, 2))) + val applied1 : Tensor1[B, Float32]= concatenated(Axis[A].at(1)) +``` + +As for split, we can also provide a tuple (or sequence) of indices, which will then select all indices in the tuple. + +```scala mdoc:silent + val selectedMultiple : Tensor2[A, B, Float32] = concatenated.select(Axis[A].at((0, 2))) +``` + +To change the values at a selection, `updated` returns a copy of the tensor with new values there, +as `updated` does for Scala's immutable collections: + +```scala mdoc:silent + val updated1 : Tensor2[A, B, Float32] = concatenated.updated(Axis[A].at(1))(Tensor(Shape(Axis[B] -> 2)).fill(0.0f)) ``` #### Squeezing, Expanding and transposing axes