Skip to content

Add commonly-used methods on tensors #164

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

Open
marcelluethi wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Add commonly-used methods on tensors #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fb87b5e
add sort, cumsum, cumprod and diff methods
marcelluethi Sep 5, 2026
92fe377
add arc..., roundingops and isnan/isfinite
marcelluethi Sep 5, 2026
0e9e213
add modulo operations
marcelluethi Sep 5, 2026
24b73aa
add nanToNum extension method
marcelluethi Sep 5, 2026
9344c83
cumsum/cumprod without label sum/factor over all elements
marcelluethi Sep 6, 2026
e15a720
fix tests
marcelluethi Sep 6, 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
17 changes: 17 additions & 0 deletions core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand All @@ -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])

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

V: IsNumber should probably be IsFloating here? Already a bug for sin, cos, tanh, which currently work for int tenors, which should not be defined. Same for new functions.

Expand Down Expand Up @@ -173,6 +181,15 @@ 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))

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

floor, ceil, and round keep being float tensors in JAX. So this implementation V -> V is correct (for IsFloating). Integer tensors are converted to float tensors so currently incorrect Int32 -> Float32, however Int32 should just not be supported (see comment above).

Probably fine as-is due to JAX's default behavior. A cast to Int seems wrong. Does a floorIntmake sense here? Probably not, as types tell the story and the user can simply write floor.asInt32. Just want this to be an active decision.

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 nanToNum: Tensor[T, V] = Tensor(Jax.jnp.nan_to_num(t.jaxValue))

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should provide arguments for nan=0.0, posinf=None, neginf=None that it passes to JAX. I would also consider removing the default for nan, so the user must explicitly specify 0.0.

https://docs.jax.dev/en/latest/_autosummary/jax.numpy.nan_to_num.html

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe like this:

def nanToNum(valueForNan: Double, valueForPosInf: Double, valueForPosNegInf: Double): Tensor[T, V]

def nanToNum(valueForNan: Double): Tensor[T, V] = nanToNum(valueForNan, valueForNan, valueForNan)
```


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] =
Expand Down
16 changes: 16 additions & 0 deletions core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would not support the default to the last axis.

@marcelluethi marcelluethi Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure about it. On one hand it is a confusing default behavior. On the other hand it gracefully handles the case of Tensor1. As far as I know, we cannot have separate extension methods with the same name on Tensor1 and generic Tensor. Argsort, argmin, etc all have this kind of Axis less version.

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I did not express my review clearly before (too early in the morning :D)

I would NOT support the default case. The user should always be explicit about which axis the sort is applied to:

val t: Tensor2[Batch, TimeStep, Int32] = ???
t.sort(Axis[TimeStep])

Even if Feature is the last dimension. DimWit works completely without positional assumptions; suddenly having a default here is wrong.

Actually, we should be more extreme and define sort only on Tensor1. Then the above statement must be:

t.vapply(Axis[TimeStep])(_.sort)
// or
t.vmap(Axis[Batch])(_.sort)
t.sort // compile-error => axis param missing

This would be identical to how linear layers or softmax work now.

val t: Tensor2[Batch, Feature, Float32] = ???
t.vmap(Axis[Batch])(linearLayer)
t.vapply(Axis[Feature])(softmax)
def softmax[L: Label, V: IsFloating](t: Tensor1[L, V]): Tensor1[L, V] =
  liftPyTensor(Jax.jnn.softmax(toPyTensor(t), axis = 0))

.sort is a function on Tensor1: Taking a vector and sorting that vector. It does not know anything about higher-dimensional tensors. This is the strict and minimal conceptual scope of .sort.

Note that some functions are more general than their minimal scope, like .dot and relu. So your version of sort wouldn't be the only one, but I think we should keep scopes very strict, especially for less common methods. With application to higher tensors with vapply and vmap.

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Additionally, sort is not a ReductionOps. Same for argsort actually...


/** 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))

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would return a flattened vector, and the type Tensor[T, V] is incorrect. See the axis comment at:
https://docs.jax.dev/en/latest/_autosummary/jax.numpy.cumsum.html

We could change the return value to Tensor1[R, V] with merger: AxesMerger.Aux; see flatten to fix this, or not support this (for now).

Actually, should cumsum just be an operation on Tensor1? Similar to sort.

t.vapply(Axis[A])(_.cumsum)

This would allow:

t.flatten.cumsum

For the default flatten case.

@benikm91 benikm91 Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Motivation is similar: cumsum is an operation over a list of values, which is a Vector / Tensor1 in tensorland.

@benikm91 benikm91 Sep 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we decide what to do here, do the same for cumprod. And diff(I think).


/** 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))

/** 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)
// ---------------------------------------------------------
Expand Down
34 changes: 34 additions & 0 deletions core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,35 @@ 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("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))
(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)))

Expand All @@ -75,6 +104,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))

Expand Down
82 changes: 82 additions & 0 deletions core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.flatten.relabelTo(Axis[A]).cumsum(Axis[A])

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.flatten.relabelTo(Axis[A]).cumprod(Axis[A])

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)
Expand Down
Loading