Skip to content

Add arange method #167

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
benikm91 merged 3 commits into from
Sep 22, 2026
Merged

Add arange method #167

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ba05194
Add arange method
benikm91 Sep 19, 2026
d3ad631
Move to scala range
benikm91 Sep 20, 2026
b5be55c
Rename arange to fromRange; add linspace
benikm91 Sep 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 92 additions & 8 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,90 @@ val notAMatrix = Tensor1(Axis[A] -> 3).eye
// ^^^^^^^^^^^^^^^^^^^^^^^^^
```

### Integer Ranges with `fromRange`

`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]).fromRange(0 until 4)

// 2, 3, 4, 5
val inclusive = Tensor1(Axis[A]).fromRange(2 to 5)

// 0, 3, 6
val stepped = Tensor1(Axis[A]).fromRange(0 until 7 by 3)

// 3, 2, 1
val descending = Tensor1(Axis[A]).fromRange(3 until 0 by -1)

// 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]).fromRange(0 until 4)
```

```scala
// ERROR: fromRange only exists on the rank 1 factory
val notAVector = Tensor2(Axis[A], Axis[B]).fromRange(0 until 4)
// error:
// 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]).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

```scala
Expand Down Expand Up @@ -390,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 58 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 104
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 117
//
```

Expand Down Expand Up @@ -435,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 58 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 113
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 126
//
```

Expand Down Expand Up @@ -528,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 135 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 138
// 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 136 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 139
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 152
//
```

Expand Down
25 changes: 25 additions & 0 deletions core/src/main/scala/dimwit/tensor/Tensor.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -279,6 +280,19 @@ 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.
* A `Range` only holds integers; for evenly spaced floating point values use [[linspace]].
*/
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]):

def fromArray(values: Array[Boolean])(using IsBoolean[V]): Tensor1[L, V] = ArrayWriter.fromArray[Tuple1[L], V](Shape1(axis -> values.length), values)
Expand All @@ -289,6 +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.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)

Expand Down
76 changes: 76 additions & 0 deletions core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,3 +129,79 @@ 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("fromRange"):

it("until: half-open interval"):
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]).fromRange(2 until 5) shouldEqual Tensor1(Axis[A]).fromArray(Array(2, 3, 4))

it("to: inclusive interval"):
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]).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]).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]).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]).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]).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
71 changes: 71 additions & 0 deletions mdocs/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,77 @@ val intEye = Tensor2(Axis[A] -> 3, Axis[B] -> 3).eye(VType[Int32])
val notAMatrix = Tensor1(Axis[A] -> 3).eye
```

### Integer Ranges with `fromRange`

`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]).fromRange(0 until 4)

// 2, 3, 4, 5
val inclusive = Tensor1(Axis[A]).fromRange(2 to 5)

// 0, 3, 6
val stepped = Tensor1(Axis[A]).fromRange(0 until 7 by 3)

// 3, 2, 1
val descending = Tensor1(Axis[A]).fromRange(3 until 0 by -1)

// 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]).fromRange(0 until 4)
```

```scala mdoc:fail
// 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

```scala mdoc:silent
Expand Down
Loading