From fb87b5efc800081f2087fc0c41604e8153710dcd Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 5 Sep 2026 06:19:26 +0200 Subject: [PATCH 1/5] add sort, cumsum, cumprod and diff methods --- .../tensor/tensorops/ReductionOps.scala | 16 ++++ .../tensor/TensorOpsReductionSuite.scala | 82 +++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala index 7481004d..3757b5e3 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala @@ -52,6 +52,22 @@ object ReductionOps: def argsort[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, Int32] = Tensor(Jax.jnp.argsort(t.jaxValue, axis = ev.index)) def argsort: Tensor[T, Int32] = Tensor(Jax.jnp.argsort(t.jaxValue)) + /** sorts the tensor `t` along the specified axis */ + def sort[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.sort(t.jaxValue, axis = ev.index)) + def sort: Tensor[T, V] = Tensor(Jax.jnp.sort(t.jaxValue)) + + /** computes the cumulative sum of the tensor `t` along the specified axis. */ + def cumsum[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.cumsum(t.jaxValue, axis = ev.index)) + def cumsum: Tensor[T, V] = Tensor(Jax.jnp.cumsum(t.jaxValue, axis = -1)) + + /** computes the cumulative product of the tensor `t` along the specified axis. */ + def cumprod[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.cumprod(t.jaxValue, axis = ev.index)) + def cumprod: Tensor[T, V] = Tensor(Jax.jnp.cumprod(t.jaxValue, axis = -1)) + + /** computes the discrete difference of the tensor `t` along the specified axis, reducing that axis' size by one. */ + def diff[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.diff(t.jaxValue, axis = ev.index)) + def diff: Tensor[T, V] = Tensor(Jax.jnp.diff(t.jaxValue)) + // --------------------------------------------------------- // IsFloat operations (IsFloat or IsInt) // --------------------------------------------------------- diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala index 9d4cdbd3..b847609f 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala @@ -168,6 +168,88 @@ class TensorOpsReductionSuite extends DimwitTest: ) ) + it("sort"): + val descendingAlongB = Tensor2(Axis[A], Axis[B]).fromArray( + Array( + Array(3.0f, 2.0f, 1.0f), + Array(6.0f, 5.0f, 4.0f) + ) + ) + descendingAlongB.sort shouldEqual Tensor2( + Axis[A], + Axis[B] + ).fromArray( + Array( + Array(1.0f, 2.0f, 3.0f), + Array(4.0f, 5.0f, 6.0f) + ) + ) + + it("sort axis A"): + val descendingAlongA = Tensor2(Axis[A], Axis[B]).fromArray( + Array( + Array(4.0f, 5.0f, 6.0f), + Array(1.0f, 2.0f, 3.0f) + ) + ) + val res = descendingAlongA.sort(axis = Axis[A]) + res shouldEqual Tensor2( + Axis[A], + Axis[B] + ).fromArray( + Array( + Array(1.0f, 2.0f, 3.0f), + Array(4.0f, 5.0f, 6.0f) + ) + ) + + it("sort axis B"): + val descendingAlongB = Tensor2(Axis[A], Axis[B]).fromArray( + Array( + Array(3.0f, 2.0f, 1.0f), + Array(6.0f, 5.0f, 4.0f) + ) + ) + val res = descendingAlongB.sort(axis = Axis[B]) + res shouldEqual Tensor2( + Axis[A], + Axis[B] + ).fromArray( + Array( + Array(1.0f, 2.0f, 3.0f), + Array(4.0f, 5.0f, 6.0f) + ) + ) + + it("cumsum"): + val res = t2.cumsum(axis = Axis[B]) + res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 3.0f, 6.0f, 4.0f, 9.0f, 15.0f)) + + it("cumsum default axis"): + t2.cumsum shouldEqual t2.cumsum(axis = Axis[B]) + + it("cumsum axis A"): + val res = t2.cumsum(axis = Axis[A]) + res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 2.0f, 3.0f, 5.0f, 7.0f, 9.0f)) + + it("cumprod"): + val res = t2.cumprod(axis = Axis[B]) + res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 2.0f, 6.0f, 4.0f, 20.0f, 120.0f)) + + it("cumprod default axis"): + t2.cumprod shouldEqual t2.cumprod(axis = Axis[B]) + + it("diff axis B"): + val res = t2.diff(axis = Axis[B]) + res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 1.0f, 1.0f, 1.0f)) + + it("diff default axis"): + t2.diff shouldEqual t2.diff(axis = Axis[B]) + + it("diff axis A"): + val res = t2.diff(axis = Axis[A]) + res shouldEqual Tensor.like(res).fromArray(Array(3.0f, 3.0f, 3.0f)) + describe("Boolean Reductions"): it("all"): b2.all shouldEqual Tensor0(false) From 92fe377687d0c5b27911f221085d2a224b76d3c8 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 5 Sep 2026 06:26:56 +0200 Subject: [PATCH 2/5] add arc..., roundingops and isnan/isfinite --- .../dimwit/tensor/tensorops/ElementWiseOps.scala | 8 ++++++++ .../tensor/TensorOpsElementwiseSuite.scala | 16 ++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala index 37c98f33..dc7e67ed 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala @@ -173,6 +173,14 @@ object ElementWiseOps: def sin: Tensor[T, V] = Tensor(Jax.jnp.sin(t.jaxValue)) def cos: Tensor[T, V] = Tensor(Jax.jnp.cos(t.jaxValue)) def tanh: Tensor[T, V] = Tensor(Jax.jnp.tanh(t.jaxValue)) + def arcsin: Tensor[T, V] = Tensor(Jax.jnp.arcsin(t.jaxValue)) + def arccos: Tensor[T, V] = Tensor(Jax.jnp.arccos(t.jaxValue)) + def arctan: Tensor[T, V] = Tensor(Jax.jnp.arctan(t.jaxValue)) + def floor: Tensor[T, V] = Tensor(Jax.jnp.floor(t.jaxValue)) + def ceil: Tensor[T, V] = Tensor(Jax.jnp.ceil(t.jaxValue)) + def round: Tensor[T, V] = Tensor(Jax.jnp.round(t.jaxValue)) + def isnan: Tensor[T, Bool] = Tensor(Jax.jnp.isnan(t.jaxValue)) + def isfinite: Tensor[T, Bool] = Tensor(Jax.jnp.isfinite(t.jaxValue)) def approxEquals(other: Tensor[T, V], tolerance: Float = 1e-6f): Tensor0[Bool] = approxElementEquals(other, tolerance).all def approxElementEquals(other: Tensor[T, V], tolerance: Float = 1e-6f): Tensor[T, Bool] = diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala index 4e43664e..583a4726 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala @@ -53,6 +53,22 @@ class TensorOpsElementwiseSuite extends DimwitTest: tZero.cos should approxEqual(Tensor.like(t2).fill(1f)) tZero.tanh should approxEqual(tZero) + it("arcsin/arccos/arctan"): + Tensor.like(t2).fill(0.5f).arcsin should approxEqual(Tensor.like(t2).fill((math.Pi / 6).toFloat), tolerance = 1e-5f) + Tensor.like(t2).fill(0.5f).arccos should approxEqual(Tensor.like(t2).fill((math.Pi / 3).toFloat), tolerance = 1e-5f) + Tensor.like(t2).fill(1.0f).arctan should approxEqual(Tensor.like(t2).fill((math.Pi / 4).toFloat), tolerance = 1e-5f) + + it("floor/ceil/round"): + val t = Tensor.like(t2).fromArray(Array(-1.5f, 0.4f, 1.5f, 2.6f)) + t.floor should approxEqual(Tensor.like(t2).fromArray(Array(-2.0f, 0.0f, 1.0f, 2.0f))) + t.ceil should approxEqual(Tensor.like(t2).fromArray(Array(-1.0f, 1.0f, 2.0f, 3.0f))) + t.round should approxEqual(Tensor.like(t2).fromArray(Array(-2.0f, 0.0f, 2.0f, 3.0f))) + + it("isnan/isfinite"): + val t = Tensor.like(t2).fromArray(Array(Float.NaN, Float.PositiveInfinity, 1.0f, 0.0f)) + t.isnan shouldEqual Tensor.like(b2).fromArray(Array(true, false, false, false)) + t.isfinite shouldEqual Tensor.like(b2).fromArray(Array(false, false, true, true)) + it("clip"): t2.clip(0.0f, 2.0f) should approxEqual(Tensor.like(t2).fromArray(Array(0.0f, 0.0f, 1.0f, 2.0f))) From 0e9e213681fc0b4e0f39d132a04d9a7bd30d69ef Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 5 Sep 2026 06:40:30 +0200 Subject: [PATCH 3/5] add modulo operations --- .../dimwit/tensor/tensorops/ElementWiseOps.scala | 8 ++++++++ .../dimwit/tensor/TensorOpsElementwiseSuite.scala | 14 ++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala index dc7e67ed..6d7d0573 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala @@ -126,12 +126,19 @@ object ElementWiseOps: /** Multiplies each element of a tensor by a scalar tensor, returning a new tensor. */ def multiplyScalar[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], s: Tensor0[V]): Tensor[T, V] = Tensor(Jax.jnp.multiply(t1.jaxValue, s.jaxValue)) + /** Computes the element-wise remainder of `t1 / t2`, matching Python's `%` operator (the result takes the sign of the divisor). */ + def mod[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.mod(t1.jaxValue, t2.jaxValue)) + + /** Computes the remainder of dividing each element of a tensor by a scalar tensor. */ + def modScalar[T <: Tuple: Labels, V: IsNumber](t1: Tensor[T, V], s: Tensor0[V]): Tensor[T, V] = Tensor(Jax.jnp.mod(t1.jaxValue, s.jaxValue)) + // extension methods for the binary operations on two tensors extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]) def +(other: Tensor[T, V]): Tensor[T, V] = add(t, other) def -(other: Tensor[T, V]): Tensor[T, V] = subtract(t, other) def *(other: Tensor[T, V]): Tensor[T, V] = multiply(t, other) + def %(other: Tensor[T, V]): Tensor[T, V] = mod(t, other) // extension methods for the scalar operations. extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]) @@ -143,6 +150,7 @@ object ElementWiseOps: def *![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(multiply) def scale(other: Tensor0[V]): Tensor[T, V] = multiplyScalar(t, other) + def %![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(mod) // extension methods extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V]) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala index 583a4726..688e9c47 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala @@ -69,6 +69,15 @@ class TensorOpsElementwiseSuite extends DimwitTest: t.isnan shouldEqual Tensor.like(b2).fromArray(Array(true, false, false, false)) t.isfinite shouldEqual Tensor.like(b2).fromArray(Array(false, false, true, true)) + it("mod"): + val t = Tensor.like(t2).fromArray(Array(-7.0f, 7.0f, -7.0f, 7.0f)) + val divisor = Tensor.like(t2).fromArray(Array(3.0f, 3.0f, -3.0f, -3.0f)) + (t % divisor) should approxEqual(Tensor.like(t2).fromArray(Array(2.0f, 1.0f, -1.0f, -2.0f))) + + it("mod broadcasting (%!)"): + val t = Tensor1(Axis[A]).fromArray(Array(-7.0f, 7.0f)) + (t %! Tensor0(3.0f)) should approxEqual(Tensor1(Axis[A]).fromArray(Array(2.0f, 1.0f))) + it("clip"): t2.clip(0.0f, 2.0f) should approxEqual(Tensor.like(t2).fromArray(Array(0.0f, 0.0f, 1.0f, 2.0f))) @@ -91,6 +100,11 @@ class TensorOpsElementwiseSuite extends DimwitTest: it("pow"): i2.pow(Tensor0(3)) shouldEqual Tensor.like(i2).fromArray(Array(-1, 0, 1, 8)) + it("mod"): + val t = Tensor.like(i2).fromArray(Array(-7, 7, -7, 7)) + val divisor = Tensor.like(i2).fromArray(Array(3, 3, -3, -3)) + (t % divisor) shouldEqual Tensor.like(i2).fromArray(Array(2, 1, -1, -2)) + it("clip"): i2.clip(0, 1) shouldEqual Tensor.like(i2).fromArray(Array(0, 0, 1, 1)) From 24b73aa5de8e5078e576e38d3e8622adcbe75cae Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 5 Sep 2026 06:42:35 +0200 Subject: [PATCH 4/5] add nanToNum extension method --- .../main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala | 1 + .../test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala index 6d7d0573..11079502 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala @@ -189,6 +189,7 @@ object ElementWiseOps: def round: Tensor[T, V] = Tensor(Jax.jnp.round(t.jaxValue)) def isnan: Tensor[T, Bool] = Tensor(Jax.jnp.isnan(t.jaxValue)) def isfinite: Tensor[T, Bool] = Tensor(Jax.jnp.isfinite(t.jaxValue)) + def nanToNum: Tensor[T, V] = Tensor(Jax.jnp.nan_to_num(t.jaxValue)) def approxEquals(other: Tensor[T, V], tolerance: Float = 1e-6f): Tensor0[Bool] = approxElementEquals(other, tolerance).all def approxElementEquals(other: Tensor[T, V], tolerance: Float = 1e-6f): Tensor[T, Bool] = diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala index 688e9c47..4b6ff62b 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala @@ -69,6 +69,10 @@ class TensorOpsElementwiseSuite extends DimwitTest: t.isnan shouldEqual Tensor.like(b2).fromArray(Array(true, false, false, false)) t.isfinite shouldEqual Tensor.like(b2).fromArray(Array(false, false, true, true)) + it("nanToNum"): + val t = Tensor.like(t2).fromArray(Array(Float.NaN, Float.PositiveInfinity, Float.NegativeInfinity, 1.0f)) + t.nanToNum shouldEqual Tensor.like(t2).fromArray(Array(0.0f, Float.MaxValue, -Float.MaxValue, 1.0f)) + it("mod"): val t = Tensor.like(t2).fromArray(Array(-7.0f, 7.0f, -7.0f, 7.0f)) val divisor = Tensor.like(t2).fromArray(Array(3.0f, 3.0f, -3.0f, -3.0f)) From 9344c830e4b8db55a748808db856a141988ead45 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sun, 6 Sep 2026 10:48:24 +0200 Subject: [PATCH 5/5] cumsum/cumprod without label sum/factor over all elements --- .../src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala index 3757b5e3..9e86dddc 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala @@ -58,11 +58,11 @@ object ReductionOps: /** computes the cumulative sum of the tensor `t` along the specified axis. */ def cumsum[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.cumsum(t.jaxValue, axis = ev.index)) - def cumsum: Tensor[T, V] = Tensor(Jax.jnp.cumsum(t.jaxValue, axis = -1)) + def cumsum: Tensor[T, V] = Tensor(Jax.jnp.cumsum(t.jaxValue)) /** computes the cumulative product of the tensor `t` along the specified axis. */ def cumprod[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.cumprod(t.jaxValue, axis = ev.index)) - def cumprod: Tensor[T, V] = Tensor(Jax.jnp.cumprod(t.jaxValue, axis = -1)) + def cumprod: Tensor[T, V] = Tensor(Jax.jnp.cumprod(t.jaxValue)) /** computes the discrete difference of the tensor `t` along the specified axis, reducing that axis' size by one. */ def diff[L: Label](axis: Axis[L])(using ev: AxisIndex[T, L]): Tensor[T, V] = Tensor(Jax.jnp.diff(t.jaxValue, axis = ev.index))