From ba051943ec3c9e449c7dc73f51b17c1d754ea36c Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Sun, 20 Sep 2026 01:00:31 +0200 Subject: [PATCH 1/3] Add arange method --- AGENTS.md | 48 +++++++++++++++---- .../src/main/scala/dimwit/tensor/Tensor.scala | 30 ++++++++++++ .../dimwit/tensor/TensorCreationSuite.scala | 32 +++++++++++++ mdocs/AGENTS.md | 27 +++++++++++ 4 files changed, 129 insertions(+), 8 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 52a46ff..ff82b4f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -197,6 +197,38 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye // ^^^^^^^^^^^^^^^^^^^^^^^^^ ``` +### Ranges with `arange` + +`arange` creates a vector of evenly spaced values in the half-open interval `[start, stop)`, like `jnp.arange`. +The extent of the axis follows from the arguments, so it is a method on the `Tensor1(Axis[L])` factory. + +```scala +// 0, 1, 2, 3 +val range = Tensor1(Axis[A]).arange(4) + +// 2, 3, 4 +val fromStart = Tensor1(Axis[A]).arange(2, 5) + +// 0, 3, 6: the extent is ceil((stop - start) / step) +val stepped = Tensor1(Axis[A]).arange(0, 7, 3) + +// Like fill and fromArray, the value type follows the arguments: Int -> Int32, Float -> Float32 +val floatRange = Tensor1(Axis[A]).arange(0.0f, 1.0f, 0.25f) + +// ... or is given explicitly on the typed factory +val shortRange = Tensor1(Axis[A], VType[Int16]).arange(4) +``` + +```scala +// ERROR: arange only exists on the rank 1 factory +val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +// error: +// value arange is not a member of dimwit.tensor.Tensor2.Axes2Factory[repl.MdocSession.MdocApp.A, +// repl.MdocSession.MdocApp.B] +// val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +// ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +``` + ### Type Aliases for Common Shapes ```scala @@ -390,10 +422,10 @@ val wrong = t.sum(Axis[C]) // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 58 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 63 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 104 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 109 // ``` @@ -435,10 +467,10 @@ val wrong = t + 5.0f // Use +! instead // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 58 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 63 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 113 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 118 // ``` @@ -528,19 +560,19 @@ val wrong = m1.dot(Axis[B])(m2) // Conflicting definitions: // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 135 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 140 and // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 138 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 143 // // error: // Conflicting definitions: // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.B, MdocApp1.this.C, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 136 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 141 and // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.C, MdocApp1.this.D, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 139 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 144 // ``` diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala index 21436b6..5eb0f6a 100644 --- a/core/src/main/scala/dimwit/tensor/Tensor.scala +++ b/core/src/main/scala/dimwit/tensor/Tensor.scala @@ -279,6 +279,22 @@ object Tensor1: def fromArray(values: Array[Float]): Tensor1[L, Float32] = Tensor1(axis, VType[Float32]).fromArray(values) def fromArray(values: Array[Double]): Tensor1[L, Float64] = Tensor1(axis, VType[Float64]).fromArray(values) + /** Creates a vector of evenly spaced values in the half-open interval `[start, stop)`, + * like `jnp.arange`. The extent of the axis is `ceil((stop - start) / step)`. + */ + def arange(stop: Int): Tensor1[L, Int32] = arange(0, stop) + def arange(start: Int, stop: Int): Tensor1[L, Int32] = arange(start, stop, 1) + def arange(start: Int, stop: Int, step: Int): Tensor1[L, Int32] = Tensor1(axis, VType[Int32]).arange(start, stop, step) + def arange(stop: Long): Tensor1[L, Int64] = arange(0L, stop) + def arange(start: Long, stop: Long): Tensor1[L, Int64] = arange(start, stop, 1L) + def arange(start: Long, stop: Long, step: Long): Tensor1[L, Int64] = Tensor1(axis, VType[Int64]).arange(start, stop, step) + def arange(stop: Float): Tensor1[L, Float32] = arange(0f, stop) + def arange(start: Float, stop: Float): Tensor1[L, Float32] = arange(start, stop, 1f) + def arange(start: Float, stop: Float, step: Float): Tensor1[L, Float32] = Tensor1(axis, VType[Float32]).arange(start, stop, step) + def arange(stop: Double): Tensor1[L, Float64] = arange(0d, stop) + def arange(start: Double, stop: Double): Tensor1[L, Float64] = arange(start, stop, 1d) + def arange(start: Double, stop: Double, step: Double): Tensor1[L, Float64] = Tensor1(axis, VType[Float64]).arange(start, stop, step) + class AxisTypedFactory[L: Label, V](axis: Axis[L], vtype: VType[V]): def fromArray(values: Array[Boolean])(using IsBoolean[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) @@ -289,6 +305,20 @@ object Tensor1: def fromArray(values: Array[Float])(using IsFloating[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) def fromArray(values: Array[Double])(using IsFloating[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) + /** @see [[AxisFactory.arange]] */ + def arange(stop: Int)(using IsInteger[V]): Tensor1[L, V] = arange(0, stop) + def arange(start: Int, stop: Int)(using IsInteger[V]): Tensor1[L, V] = arange(start, stop, 1) + def arange(start: Int, stop: Int, step: Int)(using IsInteger[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) + def arange(stop: Long)(using IsInteger[V]): Tensor1[L, V] = arange(0L, stop) + def arange(start: Long, stop: Long)(using IsInteger[V]): Tensor1[L, V] = arange(start, stop, 1L) + def arange(start: Long, stop: Long, step: Long)(using IsInteger[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) + def arange(stop: Float)(using IsFloating[V]): Tensor1[L, V] = arange(0f, stop) + def arange(start: Float, stop: Float)(using IsFloating[V]): Tensor1[L, V] = arange(start, stop, 1f) + def arange(start: Float, stop: Float, step: Float)(using IsFloating[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) + def arange(stop: Double)(using IsFloating[V]): Tensor1[L, V] = arange(0d, stop) + def arange(start: Double, stop: Double)(using IsFloating[V]): Tensor1[L, V] = arange(start, stop, 1d) + def arange(start: Double, stop: Double, step: Double)(using IsFloating[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) + def apply[L: Label](axis: Axis[L]): AxisFactory[L] = AxisFactory(axis) def apply[L: Label, V](axis: Axis[L], vtype: VType[V]): AxisTypedFactory[L, V] = AxisTypedFactory(axis, vtype) diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala index 3a3a237f..08a37cb 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -129,3 +129,35 @@ class TensorCreationSuite extends DimwitTest: Tensor2(Axis[A] -> 2, Axis[B] -> 3).eye.dtype shouldBe DType.Float32 Tensor2(Axis[A] -> 2, Axis[B] -> 3).eye(VType[Int32]).dtype shouldBe DType.Int32 Tensor2(Shape2(Axis[A] -> 2, Axis[B] -> 3)).eye(VType[Int16]).dtype shouldBe DType.Int16 + + describe("arange"): + + it("stop only: starts at 0 with step 1"): + val result = Tensor1(Axis[A]).arange(4) + result.shape shouldEqual Shape1(Axis[A] -> 4) + result shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 1, 2, 3)) + + it("start and stop: half-open interval"): + Tensor1(Axis[A]).arange(2, 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4)) + + it("step: extent is ceil((stop - start) / step), negative steps count down"): + Tensor1(Axis[A]).arange(0, 7, 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) + Tensor1(Axis[A]).arange(3, 0, -1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3, 2, 1)) + + it("empty interval gives an empty vector"): + Tensor1(Axis[A]).arange(0).shape shouldEqual Shape1(Axis[A] -> 0) + Tensor1(Axis[A]).arange(5, 2).shape shouldEqual Shape1(Axis[A] -> 0) + + it("derives the value type from the arguments"): + Tensor1(Axis[A]).arange(3).dtype shouldBe DType.Int32 + Tensor1(Axis[A]).arange(0f, 1f, 0.25f) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 0.25f, 0.5f, 0.75f)) + withJaxX64Support: + Tensor1(Axis[A]).arange(0.5).dtype shouldBe DType.Float64 + + it("typed factory takes the vtype"): + Tensor1(Axis[A], VType[Int16]).arange(3).dtype shouldBe DType.Int16 + Tensor1(Axis[A], VType[Float32]).arange(1f, 3f).dtype shouldBe DType.Float32 + + it("can be consumed as gather indices by take"): + val t = Tensor1(Axis[A]).fromArray(Array(10.0f, 20.0f, 30.0f)) + t.take(Axis[A])(Tensor1(Axis[B]).arange(3)) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 4b6e003..c0cbd4a 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -161,6 +161,33 @@ val intEye = Tensor2(Axis[A] -> 3, Axis[B] -> 3).eye(VType[Int32]) val notAMatrix = Tensor1(Axis[A] -> 3).eye ``` +### Ranges with `arange` + +`arange` creates a vector of evenly spaced values in the half-open interval `[start, stop)`, like `jnp.arange`. +The extent of the axis follows from the arguments, so it is a method on the `Tensor1(Axis[L])` factory. + +```scala mdoc:silent +// 0, 1, 2, 3 +val range = Tensor1(Axis[A]).arange(4) + +// 2, 3, 4 +val fromStart = Tensor1(Axis[A]).arange(2, 5) + +// 0, 3, 6: the extent is ceil((stop - start) / step) +val stepped = Tensor1(Axis[A]).arange(0, 7, 3) + +// Like fill and fromArray, the value type follows the arguments: Int -> Int32, Float -> Float32 +val floatRange = Tensor1(Axis[A]).arange(0.0f, 1.0f, 0.25f) + +// ... or is given explicitly on the typed factory +val shortRange = Tensor1(Axis[A], VType[Int16]).arange(4) +``` + +```scala mdoc:fail +// ERROR: arange only exists on the rank 1 factory +val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +``` + ### Type Aliases for Common Shapes ```scala mdoc:silent From d3ad6318794e9adfcfa785fb649b9ad33ec47b24 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Sun, 20 Sep 2026 15:58:20 +0200 Subject: [PATCH 2/3] Move to scala range --- AGENTS.md | 46 ++++++++++--------- .../src/main/scala/dimwit/tensor/Tensor.scala | 36 ++++----------- .../dimwit/tensor/TensorCreationSuite.scala | 40 ++++++++-------- mdocs/AGENTS.md | 28 ++++++----- 4 files changed, 72 insertions(+), 78 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ff82b4f..9ae40df 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -199,33 +199,37 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye ### Ranges with `arange` -`arange` creates a vector of evenly spaced values in the half-open interval `[start, stop)`, like `jnp.arange`. -The extent of the axis follows from the arguments, so it is a method on the `Tensor1(Axis[L])` factory. +`arange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. +The extent of the axis is the length of the range, so it is a method on the `Tensor1(Axis[L])` factory. ```scala // 0, 1, 2, 3 -val range = Tensor1(Axis[A]).arange(4) +val range = Tensor1(Axis[A]).arange(0 until 4) -// 2, 3, 4 -val fromStart = Tensor1(Axis[A]).arange(2, 5) +// 2, 3, 4, 5 +val inclusive = Tensor1(Axis[A]).arange(2 to 5) -// 0, 3, 6: the extent is ceil((stop - start) / step) -val stepped = Tensor1(Axis[A]).arange(0, 7, 3) +// 0, 3, 6 +val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3) -// Like fill and fromArray, the value type follows the arguments: Int -> Int32, Float -> Float32 -val floatRange = Tensor1(Axis[A]).arange(0.0f, 1.0f, 0.25f) +// 3, 2, 1 +val descending = Tensor1(Axis[A]).arange(3 until 0 by -1) -// ... or is given explicitly on the typed factory -val shortRange = Tensor1(Axis[A], VType[Int16]).arange(4) +// A Range has no value type to derive from, so arange defaults to Int32 +// and takes the value type as an argument ... +val floatRange = Tensor1(Axis[A]).arange(0 until 4, VType[Float32]) + +// ... or from the typed factory +val shortRange = Tensor1(Axis[A], VType[Int16]).arange(0 until 4) ``` ```scala // ERROR: arange only exists on the rank 1 factory -val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) // error: // value arange is not a member of dimwit.tensor.Tensor2.Axes2Factory[repl.MdocSession.MdocApp.A, // repl.MdocSession.MdocApp.B] -// val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +// val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ``` @@ -422,10 +426,10 @@ val wrong = t.sum(Axis[C]) // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 63 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 64 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 109 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 110 // ``` @@ -467,10 +471,10 @@ val wrong = t + 5.0f // Use +! instead // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 63 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 64 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 118 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 119 // ``` @@ -560,19 +564,19 @@ val wrong = m1.dot(Axis[B])(m2) // Conflicting definitions: // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 140 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 141 and // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 143 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 144 // // error: // Conflicting definitions: // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.B, MdocApp1.this.C, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 141 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 142 and // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.C, MdocApp1.this.D, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 144 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 145 // ``` diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala index 5eb0f6a..318a868 100644 --- a/core/src/main/scala/dimwit/tensor/Tensor.scala +++ b/core/src/main/scala/dimwit/tensor/Tensor.scala @@ -10,6 +10,7 @@ import dimwit.tensor.Labels import dimwit.tensor.TensorOps.IsBoolean import dimwit.tensor.TensorOps.IsFloating import dimwit.tensor.TensorOps.IsInteger +import dimwit.tensor.TensorOps.IsNumber import dimwit.tensor.TypedIndex import dimwit.tensor.VType import me.shadaj.scalapy.py @@ -279,21 +280,11 @@ object Tensor1: def fromArray(values: Array[Float]): Tensor1[L, Float32] = Tensor1(axis, VType[Float32]).fromArray(values) def fromArray(values: Array[Double]): Tensor1[L, Float64] = Tensor1(axis, VType[Float64]).fromArray(values) - /** Creates a vector of evenly spaced values in the half-open interval `[start, stop)`, - * like `jnp.arange`. The extent of the axis is `ceil((stop - start) / step)`. + /** Creates a vector with the elements of the given range, like `jnp.arange`. + * The extent of the axis is `range.length`. */ - def arange(stop: Int): Tensor1[L, Int32] = arange(0, stop) - def arange(start: Int, stop: Int): Tensor1[L, Int32] = arange(start, stop, 1) - def arange(start: Int, stop: Int, step: Int): Tensor1[L, Int32] = Tensor1(axis, VType[Int32]).arange(start, stop, step) - def arange(stop: Long): Tensor1[L, Int64] = arange(0L, stop) - def arange(start: Long, stop: Long): Tensor1[L, Int64] = arange(start, stop, 1L) - def arange(start: Long, stop: Long, step: Long): Tensor1[L, Int64] = Tensor1(axis, VType[Int64]).arange(start, stop, step) - def arange(stop: Float): Tensor1[L, Float32] = arange(0f, stop) - def arange(start: Float, stop: Float): Tensor1[L, Float32] = arange(start, stop, 1f) - def arange(start: Float, stop: Float, step: Float): Tensor1[L, Float32] = Tensor1(axis, VType[Float32]).arange(start, stop, step) - def arange(stop: Double): Tensor1[L, Float64] = arange(0d, stop) - def arange(start: Double, stop: Double): Tensor1[L, Float64] = arange(start, stop, 1d) - def arange(start: Double, stop: Double, step: Double): Tensor1[L, Float64] = Tensor1(axis, VType[Float64]).arange(start, stop, step) + def arange[V: IsNumber](range: Range, vtype: VType[V] = VType[Int32]): Tensor1[L, V] = + Tensor1(axis, vtype).arange(range) class AxisTypedFactory[L: Label, V](axis: Axis[L], vtype: VType[V]): @@ -306,18 +297,11 @@ object Tensor1: def fromArray(values: Array[Double])(using IsFloating[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) /** @see [[AxisFactory.arange]] */ - def arange(stop: Int)(using IsInteger[V]): Tensor1[L, V] = arange(0, stop) - def arange(start: Int, stop: Int)(using IsInteger[V]): Tensor1[L, V] = arange(start, stop, 1) - def arange(start: Int, stop: Int, step: Int)(using IsInteger[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) - def arange(stop: Long)(using IsInteger[V]): Tensor1[L, V] = arange(0L, stop) - def arange(start: Long, stop: Long)(using IsInteger[V]): Tensor1[L, V] = arange(start, stop, 1L) - def arange(start: Long, stop: Long, step: Long)(using IsInteger[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) - def arange(stop: Float)(using IsFloating[V]): Tensor1[L, V] = arange(0f, stop) - def arange(start: Float, stop: Float)(using IsFloating[V]): Tensor1[L, V] = arange(start, stop, 1f) - def arange(start: Float, stop: Float, step: Float)(using IsFloating[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) - def arange(stop: Double)(using IsFloating[V]): Tensor1[L, V] = arange(0d, stop) - def arange(start: Double, stop: Double)(using IsFloating[V]): Tensor1[L, V] = arange(start, stop, 1d) - def arange(start: Double, stop: Double, step: Double)(using IsFloating[V]): Tensor1[L, V] = Tensor(Jax.jnp.arange(start, stop, step, dtype = vtype.dtype.jaxType)) + def arange(range: Range)(using IsNumber[V]): Tensor1[L, V] = + val stop = range match + case r: Range.Inclusive => r.end + r.step.sign + case r: Range.Exclusive => r.end + Tensor(Jax.jnp.arange(range.start, stop, range.step, dtype = vtype.dtype.jaxType)) def apply[L: Label](axis: Axis[L]): AxisFactory[L] = AxisFactory(axis) def apply[L: Label, V](axis: Axis[L], vtype: VType[V]): AxisTypedFactory[L, V] = AxisTypedFactory(axis, vtype) diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala index 08a37cb..50afd56 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -132,32 +132,34 @@ class TensorCreationSuite extends DimwitTest: describe("arange"): - it("stop only: starts at 0 with step 1"): - val result = Tensor1(Axis[A]).arange(4) + it("until: half-open interval"): + val result = Tensor1(Axis[A]).arange(0 until 4) result.shape shouldEqual Shape1(Axis[A] -> 4) result shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 1, 2, 3)) + Tensor1(Axis[A]).arange(2 until 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4)) - it("start and stop: half-open interval"): - Tensor1(Axis[A]).arange(2, 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4)) + it("to: inclusive interval"): + Tensor1(Axis[A]).arange(2 to 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4, 5)) + Tensor1(Axis[A]).arange(0 to 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) - it("step: extent is ceil((stop - start) / step), negative steps count down"): - Tensor1(Axis[A]).arange(0, 7, 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) - Tensor1(Axis[A]).arange(3, 0, -1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3, 2, 1)) + it("by: stepped and negative steps count down"): + Tensor1(Axis[A]).arange(0 until 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) + Tensor1(Axis[A]).arange(3 until 0 by -1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3, 2, 1)) + Tensor1(Axis[A]).arange(10 to 0 by -3) shouldEqual Tensor1(Axis[A]).fromArray(Array(10, 7, 4, 1)) - it("empty interval gives an empty vector"): - Tensor1(Axis[A]).arange(0).shape shouldEqual Shape1(Axis[A] -> 0) - Tensor1(Axis[A]).arange(5, 2).shape shouldEqual Shape1(Axis[A] -> 0) + it("empty range gives an empty vector"): + Tensor1(Axis[A]).arange(0 until 0).shape shouldEqual Shape1(Axis[A] -> 0) + Tensor1(Axis[A]).arange(5 until 2).shape shouldEqual Shape1(Axis[A] -> 0) - it("derives the value type from the arguments"): - Tensor1(Axis[A]).arange(3).dtype shouldBe DType.Int32 - Tensor1(Axis[A]).arange(0f, 1f, 0.25f) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 0.25f, 0.5f, 0.75f)) - withJaxX64Support: - Tensor1(Axis[A]).arange(0.5).dtype shouldBe DType.Float64 + it("defaults to Int32 and takes the vtype as an argument"): + Tensor1(Axis[A]).arange(0 until 3).dtype shouldBe DType.Int32 + Tensor1(Axis[A]).arange(0 until 3, VType[Int16]).dtype shouldBe DType.Int16 + Tensor1(Axis[A]).arange(0 until 3, VType[Float32]) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f, 2.0f)) - it("typed factory takes the vtype"): - Tensor1(Axis[A], VType[Int16]).arange(3).dtype shouldBe DType.Int16 - Tensor1(Axis[A], VType[Float32]).arange(1f, 3f).dtype shouldBe DType.Float32 + it("typed factory uses its vtype"): + Tensor1(Axis[A], VType[Int16]).arange(0 until 3).dtype shouldBe DType.Int16 + Tensor1(Axis[A], VType[Float32]).arange(1 until 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) it("can be consumed as gather indices by take"): val t = Tensor1(Axis[A]).fromArray(Array(10.0f, 20.0f, 30.0f)) - t.take(Axis[A])(Tensor1(Axis[B]).arange(3)) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) + t.take(Axis[A])(Tensor1(Axis[B]).arange(0 until 3)) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f)) diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index c0cbd4a..f441226 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -163,29 +163,33 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye ### Ranges with `arange` -`arange` creates a vector of evenly spaced values in the half-open interval `[start, stop)`, like `jnp.arange`. -The extent of the axis follows from the arguments, so it is a method on the `Tensor1(Axis[L])` factory. +`arange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. +The extent of the axis is the length of the range, so it is a method on the `Tensor1(Axis[L])` factory. ```scala mdoc:silent // 0, 1, 2, 3 -val range = Tensor1(Axis[A]).arange(4) +val range = Tensor1(Axis[A]).arange(0 until 4) -// 2, 3, 4 -val fromStart = Tensor1(Axis[A]).arange(2, 5) +// 2, 3, 4, 5 +val inclusive = Tensor1(Axis[A]).arange(2 to 5) -// 0, 3, 6: the extent is ceil((stop - start) / step) -val stepped = Tensor1(Axis[A]).arange(0, 7, 3) +// 0, 3, 6 +val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3) -// Like fill and fromArray, the value type follows the arguments: Int -> Int32, Float -> Float32 -val floatRange = Tensor1(Axis[A]).arange(0.0f, 1.0f, 0.25f) +// 3, 2, 1 +val descending = Tensor1(Axis[A]).arange(3 until 0 by -1) -// ... or is given explicitly on the typed factory -val shortRange = Tensor1(Axis[A], VType[Int16]).arange(4) +// A Range has no value type to derive from, so arange defaults to Int32 +// and takes the value type as an argument ... +val floatRange = Tensor1(Axis[A]).arange(0 until 4, VType[Float32]) + +// ... or from the typed factory +val shortRange = Tensor1(Axis[A], VType[Int16]).arange(0 until 4) ``` ```scala mdoc:fail // ERROR: arange only exists on the rank 1 factory -val notAVector = Tensor2(Axis[A], Axis[B]).arange(4) +val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) ``` ### Type Aliases for Common Shapes From b5be55c9a3c7e035f8f663a4f90b66759bf534d5 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Mon, 21 Sep 2026 14:23:40 +0200 Subject: [PATCH 3/3] Rename arange to fromRange; add linspace --- AGENTS.md | 94 ++++++++++++++----- .../src/main/scala/dimwit/tensor/Tensor.scala | 23 +++-- .../dimwit/tensor/TensorCreationSuite.scala | 74 +++++++++++---- mdocs/AGENTS.md | 64 ++++++++++--- 4 files changed, 198 insertions(+), 57 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 9ae40df..2c82bdd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -197,40 +197,88 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye // ^^^^^^^^^^^^^^^^^^^^^^^^^ ``` -### Ranges with `arange` +### Integer Ranges with `fromRange` -`arange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. +`fromRange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. The extent of the axis is the length of the range, so it is a method on the `Tensor1(Axis[L])` factory. ```scala // 0, 1, 2, 3 -val range = Tensor1(Axis[A]).arange(0 until 4) +val range = Tensor1(Axis[A]).fromRange(0 until 4) // 2, 3, 4, 5 -val inclusive = Tensor1(Axis[A]).arange(2 to 5) +val inclusive = Tensor1(Axis[A]).fromRange(2 to 5) // 0, 3, 6 -val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3) +val stepped = Tensor1(Axis[A]).fromRange(0 until 7 by 3) // 3, 2, 1 -val descending = Tensor1(Axis[A]).arange(3 until 0 by -1) +val descending = Tensor1(Axis[A]).fromRange(3 until 0 by -1) -// A Range has no value type to derive from, so arange defaults to Int32 -// and takes the value type as an argument ... -val floatRange = Tensor1(Axis[A]).arange(0 until 4, VType[Float32]) +// A Range has no value type to derive from, so fromRange defaults to Int32 +// and takes the (integer) value type as an argument ... +val byteRange = Tensor1(Axis[A]).fromRange(0 until 4, VType[Int8]) // ... or from the typed factory -val shortRange = Tensor1(Axis[A], VType[Int16]).arange(0 until 4) +val shortRange = Tensor1(Axis[A], VType[Int16]).fromRange(0 until 4) ``` ```scala -// ERROR: arange only exists on the rank 1 factory -val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) +// ERROR: fromRange only exists on the rank 1 factory +val notAVector = Tensor2(Axis[A], Axis[B]).fromRange(0 until 4) // error: -// value arange is not a member of dimwit.tensor.Tensor2.Axes2Factory[repl.MdocSession.MdocApp.A, +// value fromRange is not a member of dimwit.tensor.Tensor2.Axes2Factory[repl.MdocSession.MdocApp.A, // repl.MdocSession.MdocApp.B] -// val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) -// ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +// val notAVector = Tensor2(Axis[A], Axis[B]).fromRange(0 until 4) +// ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +``` + +```scala +// ERROR: a Range holds integers, so the value type must be an integer type +val floatRange = Tensor1(Axis[A]).fromRange(0 until 4, VType[Float32]) +// error: +// Operation only valid for Integer tensors. +// val floatRange = Tensor1(Axis[A]).fromRange(0 until 4, VType[Float32]) +// ^ +``` + +**Note**: A Scala `Range` only holds integers. There is deliberately no `arange` with a floating-point step, +because accumulating a step like `0.1f` is imprecise (see the `numpy.arange` docs). Use `linspace` instead. + +### Evenly Spaced Values with `linspace` + +`linspace` creates a vector of `num` evenly spaced values from `start` to `stop`, like `jnp.linspace`. +The extent of the axis is `num`, so it is a method on the `Tensor1(Axis[L])` factory. +As for other computation values (`clip`, `pow`, `learningRate`, ...), `start` and `stop` are `Tensor0`s +and determine the value type. They may be traced, e.g. `x.min` and `x.max`; only `num` must be static. + +```scala +// 0.0, 0.25, 0.5, 0.75, 1.0 +val spaced = Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 5) + +// endpoint = false excludes stop: 0.0, 0.25, 0.5, 0.75 +val halfOpen = Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 4, endpoint = false) + +// start > stop counts down: 1.0, 0.5, 0.0 +val descendingSpaced = Tensor1(Axis[A]).linspace(Tensor0(1.0f), Tensor0(0.0f), 3) + +// Data-dependent bounds, e.g. histogram bin edges +val samples = Tensor1(Axis[B]).fromArray(Array(4.0f, 2.0f, 8.0f)) +val binEdges = Tensor1(Axis[A]).linspace(samples.min, samples.max, 4) + +// The typed factory fixes the value type; with dimwit.Conversions.given +// plain literals are converted to Tensor0 of that type +import dimwit.Conversions.given +val halfSpaced = Tensor1(Axis[A], VType[Float16]).linspace(0.0f, 1.0f, 5) +``` + +```scala +// ERROR: linspace produces floating point values, so start and stop must be floating +val intSpaced = Tensor1(Axis[A]).linspace(Tensor0(0), Tensor0(1), 5) +// error: +// Operation only valid for Floating tensors. +// val intSpaced = Tensor1(Axis[A]).linspace(Tensor0(0), Tensor0(1), 5) +// ^ ``` ### Type Aliases for Common Shapes @@ -426,10 +474,10 @@ val wrong = t.sum(Axis[C]) // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 64 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 71 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 110 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 117 // ``` @@ -471,10 +519,10 @@ val wrong = t + 5.0f // Use +! instead // Conflicting definitions: // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 64 and +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 71 and // val t: // dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp0 at line 119 +// dimwit.tensor.DType.Float32] in class MdocApp0 at line 126 // ``` @@ -564,19 +612,19 @@ val wrong = m1.dot(Axis[B])(m2) // Conflicting definitions: // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 141 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 148 and // val m1: // dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 144 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 151 // // error: // Conflicting definitions: // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.B, MdocApp1.this.C, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 142 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 149 and // val m2: // dimwit.tensor.Tensor2[MdocApp1.this.C, MdocApp1.this.D, -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 145 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 152 // ``` diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala index 318a868..a2bad83 100644 --- a/core/src/main/scala/dimwit/tensor/Tensor.scala +++ b/core/src/main/scala/dimwit/tensor/Tensor.scala @@ -280,11 +280,18 @@ object Tensor1: def fromArray(values: Array[Float]): Tensor1[L, Float32] = Tensor1(axis, VType[Float32]).fromArray(values) def fromArray(values: Array[Double]): Tensor1[L, Float64] = Tensor1(axis, VType[Float64]).fromArray(values) - /** Creates a vector with the elements of the given range, like `jnp.arange`. - * The extent of the axis is `range.length`. + /** Creates a vector with the elements of the given range. + * A `Range` only holds integers; for evenly spaced floating point values use [[linspace]]. */ - def arange[V: IsNumber](range: Range, vtype: VType[V] = VType[Int32]): Tensor1[L, V] = - Tensor1(axis, vtype).arange(range) + def fromRange[V: IsInteger](range: Range, vtype: VType[V] = VType[Int32]): Tensor1[L, V] = + Tensor1(axis, vtype).fromRange(range) + + /** Creates a vector of `num` evenly spaced values over `[start, stop]`. + * With `endpoint = false` the interval is half-open `[start, stop)`, i.e. `stop` is excluded. + * `start` and `stop` may be traced (e.g. `x.min`, `x.max`), only `num` must be static. + */ + def linspace[V: IsFloating](start: Tensor0[V], stop: Tensor0[V], num: Int, endpoint: Boolean = true): Tensor1[L, V] = + Tensor1(axis, VType[V]).linspace(start, stop, num, endpoint) class AxisTypedFactory[L: Label, V](axis: Axis[L], vtype: VType[V]): @@ -296,13 +303,17 @@ object Tensor1: def fromArray(values: Array[Float])(using IsFloating[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) def fromArray(values: Array[Double])(using IsFloating[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values) - /** @see [[AxisFactory.arange]] */ - def arange(range: Range)(using IsNumber[V]): Tensor1[L, V] = + /** @see [[AxisFactory.fromRange]] */ + def fromRange(range: Range)(using IsInteger[V]): Tensor1[L, V] = val stop = range match case r: Range.Inclusive => r.end + r.step.sign case r: Range.Exclusive => r.end Tensor(Jax.jnp.arange(range.start, stop, range.step, dtype = vtype.dtype.jaxType)) + /** @see [[AxisFactory.linspace]] */ + def linspace(start: Tensor0[V], stop: Tensor0[V], num: Int, endpoint: Boolean = true)(using IsFloating[V]): Tensor1[L, V] = + Tensor(Jax.jnp.linspace(start.jaxValue, stop.jaxValue, num, endpoint = endpoint, dtype = vtype.dtype.jaxType)) + def apply[L: Label](axis: Axis[L]): AxisFactory[L] = AxisFactory(axis) def apply[L: Label, V](axis: Axis[L], vtype: VType[V]): AxisTypedFactory[L, V] = AxisTypedFactory(axis, vtype) diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala index 50afd56..1af10d8 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -130,36 +130,78 @@ class TensorCreationSuite extends DimwitTest: Tensor2(Axis[A] -> 2, Axis[B] -> 3).eye(VType[Int32]).dtype shouldBe DType.Int32 Tensor2(Shape2(Axis[A] -> 2, Axis[B] -> 3)).eye(VType[Int16]).dtype shouldBe DType.Int16 - describe("arange"): + describe("fromRange"): it("until: half-open interval"): - val result = Tensor1(Axis[A]).arange(0 until 4) + val result = Tensor1(Axis[A]).fromRange(0 until 4) result.shape shouldEqual Shape1(Axis[A] -> 4) result shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 1, 2, 3)) - Tensor1(Axis[A]).arange(2 until 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4)) + Tensor1(Axis[A]).fromRange(2 until 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4)) it("to: inclusive interval"): - Tensor1(Axis[A]).arange(2 to 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4, 5)) - Tensor1(Axis[A]).arange(0 to 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) + Tensor1(Axis[A]).fromRange(2 to 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4, 5)) + Tensor1(Axis[A]).fromRange(0 to 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) it("by: stepped and negative steps count down"): - Tensor1(Axis[A]).arange(0 until 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) - Tensor1(Axis[A]).arange(3 until 0 by -1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3, 2, 1)) - Tensor1(Axis[A]).arange(10 to 0 by -3) shouldEqual Tensor1(Axis[A]).fromArray(Array(10, 7, 4, 1)) + Tensor1(Axis[A]).fromRange(0 until 7 by 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 3, 6)) + Tensor1(Axis[A]).fromRange(3 until 0 by -1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3, 2, 1)) + Tensor1(Axis[A]).fromRange(10 to 0 by -3) shouldEqual Tensor1(Axis[A]).fromArray(Array(10, 7, 4, 1)) it("empty range gives an empty vector"): - Tensor1(Axis[A]).arange(0 until 0).shape shouldEqual Shape1(Axis[A] -> 0) - Tensor1(Axis[A]).arange(5 until 2).shape shouldEqual Shape1(Axis[A] -> 0) + Tensor1(Axis[A]).fromRange(0 until 0).shape shouldEqual Shape1(Axis[A] -> 0) + Tensor1(Axis[A]).fromRange(5 until 2).shape shouldEqual Shape1(Axis[A] -> 0) it("defaults to Int32 and takes the vtype as an argument"): - Tensor1(Axis[A]).arange(0 until 3).dtype shouldBe DType.Int32 - Tensor1(Axis[A]).arange(0 until 3, VType[Int16]).dtype shouldBe DType.Int16 - Tensor1(Axis[A]).arange(0 until 3, VType[Float32]) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f, 2.0f)) + Tensor1(Axis[A]).fromRange(0 until 3).dtype shouldBe DType.Int32 + Tensor1(Axis[A]).fromRange(0 until 3, VType[Int16]).dtype shouldBe DType.Int16 + Tensor1(Axis[A]).fromRange(0 until 3, VType[Int16]).asInt32 shouldEqual Tensor1(Axis[A]).fromArray(Array(0, 1, 2)) it("typed factory uses its vtype"): - Tensor1(Axis[A], VType[Int16]).arange(0 until 3).dtype shouldBe DType.Int16 - Tensor1(Axis[A], VType[Float32]).arange(1 until 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + Tensor1(Axis[A], VType[Int16]).fromRange(0 until 3).dtype shouldBe DType.Int16 + Tensor1(Axis[A], VType[Int16]).fromRange(1 until 3).asInt32 shouldEqual Tensor1(Axis[A]).fromArray(Array(1, 2)) + + it("rejects non-integer vtypes at compile time"): + 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"): val t = Tensor1(Axis[A]).fromArray(Array(10.0f, 20.0f, 30.0f)) - t.take(Axis[A])(Tensor1(Axis[B]).arange(0 until 3)) shouldEqual Tensor1(Axis[B]).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)) + + describe("linspace"): + + it("num evenly spaced values including the endpoint"): + val result = Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 5) + result.shape shouldEqual Shape1(Axis[A] -> 5) + result shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 0.25f, 0.5f, 0.75f, 1.0f)) + Tensor1(Axis[A]).linspace(Tensor0(2.0f), Tensor0(3.0f), 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(2.0f, 2.5f, 3.0f)) + + it("endpoint = false excludes stop"): + Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 4, endpoint = false) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 0.25f, 0.5f, 0.75f)) + + it("descending when start > stop"): + Tensor1(Axis[A]).linspace(Tensor0(1.0f), Tensor0(0.0f), 3) shouldEqual Tensor1(Axis[A]).fromArray(Array(1.0f, 0.5f, 0.0f)) + + it("num = 1 gives start, num = 0 gives an empty vector"): + Tensor1(Axis[A]).linspace(Tensor0(3.0f), Tensor0(7.0f), 1) shouldEqual Tensor1(Axis[A]).fromArray(Array(3.0f)) + Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 0).shape shouldEqual Shape1(Axis[A] -> 0) + + it("value type is that of start and stop"): + Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 3).dtype shouldBe DType.Float32 + Tensor1(Axis[A]).linspace(Tensor0(VType[Float16])(0.0f), Tensor0(VType[Float16])(1.0f), 3).dtype shouldBe DType.Float16 + withJaxX64Support: + Tensor1(Axis[A]).linspace(Tensor0(0.0), Tensor0(1.0), 3).dtype shouldBe DType.Float64 + + it("typed factory fixes the value type and accepts converted literals"): + import dimwit.Conversions.given + Tensor1(Axis[A], VType[Float16]).linspace(0.0f, 1.0f, 3).dtype shouldBe DType.Float16 + Tensor1(Axis[A], VType[Float32]).linspace(0.0f, 1.0f, 3, endpoint = false) shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f / 3.0f, 2.0f / 3.0f)) + + it("start and stop can be data dependent"): + val x = Tensor1(Axis[B]).fromArray(Array(4.0f, 2.0f, 8.0f)) + Tensor1(Axis[A]).linspace(x.min, x.max, 4) shouldEqual Tensor1(Axis[A]).fromArray(Array(2.0f, 4.0f, 6.0f, 8.0f)) + + it("rejects non-floating and mixed value types at compile time"): + typeCheckErrors("Tensor1(Axis[A]).linspace(Tensor0(0), Tensor0(3), 3)") should not be empty + typeCheckErrors("Tensor1(Axis[A], VType[Int32]).linspace(Tensor0(0), Tensor0(3), 3)") should not be empty + typeCheckErrors("Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0), 3)") should not be empty diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index f441226..e43be0d 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -161,35 +161,75 @@ val intEye = Tensor2(Axis[A] -> 3, Axis[B] -> 3).eye(VType[Int32]) val notAMatrix = Tensor1(Axis[A] -> 3).eye ``` -### Ranges with `arange` +### Integer Ranges with `fromRange` -`arange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. +`fromRange` creates a vector with the elements of a Scala `Range`, like `jnp.arange`. The extent of the axis is the length of the range, so it is a method on the `Tensor1(Axis[L])` factory. ```scala mdoc:silent // 0, 1, 2, 3 -val range = Tensor1(Axis[A]).arange(0 until 4) +val range = Tensor1(Axis[A]).fromRange(0 until 4) // 2, 3, 4, 5 -val inclusive = Tensor1(Axis[A]).arange(2 to 5) +val inclusive = Tensor1(Axis[A]).fromRange(2 to 5) // 0, 3, 6 -val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3) +val stepped = Tensor1(Axis[A]).fromRange(0 until 7 by 3) // 3, 2, 1 -val descending = Tensor1(Axis[A]).arange(3 until 0 by -1) +val descending = Tensor1(Axis[A]).fromRange(3 until 0 by -1) -// A Range has no value type to derive from, so arange defaults to Int32 -// and takes the value type as an argument ... -val floatRange = Tensor1(Axis[A]).arange(0 until 4, VType[Float32]) +// A Range has no value type to derive from, so fromRange defaults to Int32 +// and takes the (integer) value type as an argument ... +val byteRange = Tensor1(Axis[A]).fromRange(0 until 4, VType[Int8]) // ... or from the typed factory -val shortRange = Tensor1(Axis[A], VType[Int16]).arange(0 until 4) +val shortRange = Tensor1(Axis[A], VType[Int16]).fromRange(0 until 4) ``` ```scala mdoc:fail -// ERROR: arange only exists on the rank 1 factory -val notAVector = Tensor2(Axis[A], Axis[B]).arange(0 until 4) +// ERROR: fromRange only exists on the rank 1 factory +val notAVector = Tensor2(Axis[A], Axis[B]).fromRange(0 until 4) +``` + +```scala mdoc:fail +// ERROR: a Range holds integers, so the value type must be an integer type +val floatRange = Tensor1(Axis[A]).fromRange(0 until 4, VType[Float32]) +``` + +**Note**: A Scala `Range` only holds integers. There is deliberately no `arange` with a floating-point step, +because accumulating a step like `0.1f` is imprecise (see the `numpy.arange` docs). Use `linspace` instead. + +### Evenly Spaced Values with `linspace` + +`linspace` creates a vector of `num` evenly spaced values from `start` to `stop`, like `jnp.linspace`. +The extent of the axis is `num`, so it is a method on the `Tensor1(Axis[L])` factory. +As for other computation values (`clip`, `pow`, `learningRate`, ...), `start` and `stop` are `Tensor0`s +and determine the value type. They may be traced, e.g. `x.min` and `x.max`; only `num` must be static. + +```scala mdoc:silent +// 0.0, 0.25, 0.5, 0.75, 1.0 +val spaced = Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 5) + +// endpoint = false excludes stop: 0.0, 0.25, 0.5, 0.75 +val halfOpen = Tensor1(Axis[A]).linspace(Tensor0(0.0f), Tensor0(1.0f), 4, endpoint = false) + +// start > stop counts down: 1.0, 0.5, 0.0 +val descendingSpaced = Tensor1(Axis[A]).linspace(Tensor0(1.0f), Tensor0(0.0f), 3) + +// Data-dependent bounds, e.g. histogram bin edges +val samples = Tensor1(Axis[B]).fromArray(Array(4.0f, 2.0f, 8.0f)) +val binEdges = Tensor1(Axis[A]).linspace(samples.min, samples.max, 4) + +// The typed factory fixes the value type; with dimwit.Conversions.given +// plain literals are converted to Tensor0 of that type +import dimwit.Conversions.given +val halfSpaced = Tensor1(Axis[A], VType[Float16]).linspace(0.0f, 1.0f, 5) +``` + +```scala mdoc:fail +// ERROR: linspace produces floating point values, so start and stop must be floating +val intSpaced = Tensor1(Axis[A]).linspace(Tensor0(0), Tensor0(1), 5) ``` ### Type Aliases for Common Shapes