Add arange method

AGENTS.md

@@ -197,6 +197,42 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye
 //                  ^^^^^^^^^^^^^^^^^^^^^^^^^
 ```
 
+### Ranges with `arange`
+
+`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(0 until 4)
+
+// 2, 3, 4, 5
+val inclusive = Tensor1(Axis[A]).arange(2 to 5)
+
+// 0, 3, 6
+val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3)
+
+// 3, 2, 1
+val descending = Tensor1(Axis[A]).arange(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])
+
+// ... 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(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(0 until 4)
+//                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+```
+
 ### Type Aliases for Common Shapes
 
 ```scala
@@ -390,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 58 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 104
+//     dimwit.tensor.DType.Float32] in class MdocApp0 at line 110
 //
 ```
 
@@ -435,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 58 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 113
+//     dimwit.tensor.DType.Float32] in class MdocApp0 at line 119
 //
 ```
 
@@ -528,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 135 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 138
+//     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 136 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 139
+//     dimwit.tensor.DType.Float32] in class MdocApp1 at line 145
 //
 ```
 

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,6 +280,12 @@ 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`.
+      */
+    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]):
 
     def fromArray(values: Array[Boolean])(using IsBoolean[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values)
@@ -289,6 +296,13 @@ 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] =
+      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)
 

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

@@ -129,3 +129,37 @@ 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("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("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("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 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("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 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(0 until 3)) shouldEqual Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f, 30.0f))

mdocs/AGENTS.md

@@ -161,6 +161,37 @@ 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 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)
+
+// 2, 3, 4, 5
+val inclusive = Tensor1(Axis[A]).arange(2 to 5)
+
+// 0, 3, 6
+val stepped = Tensor1(Axis[A]).arange(0 until 7 by 3)
+
+// 3, 2, 1
+val descending = Tensor1(Axis[A]).arange(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])
+
+// ... 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(0 until 4)
+```
+
 ### Type Aliases for Common Shapes
 
 ```scala mdoc:silent