Add commonly-used methods on tensors by marcelluethi · Pull Request #164 · dimwit-dev/dimwit

marcelluethi · 2026-09-05T04:50:07Z

DimWits wrapping of the basic tensor methods supported in jax is rather patchy and was introduced on a per need basis.
This PR proposes to add some more, frequently used tensor operations. Each of them is just a one-liner, wrapping the underlying jax function.

The methods added are:

sort
cumsum, cumprod
diff
floor, ceil, round
arcsin, arccos, arctan
isnan, isfinite, nanToNum
mod (with a % operator)

benikm91 · 2026-09-06T07:03:33Z


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


I would not support the default to the last axis.

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.

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.

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

benikm91 · 2026-09-06T07:05:34Z

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


I would not support the default to the last axis.

Good catch. The default should not be axis=-1 but rather just `Tensor(Jax.jnp.cumsum(t.jaxValue) if we want to be consistent. Fixed in the last commit

benikm91 · 2026-09-06T14:48:26Z

+    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])


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.

benikm91 · 2026-09-06T14:52:18Z

+    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))


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.

benikm91 · 2026-09-06T14:54:11Z

+    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))


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

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) ```

benikm91 · 2026-09-06T15:05:55Z

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


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.

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

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

benikm91 · 2026-09-06T15:16:06Z

@marcelluethi I made some targeted comments for specific lines of code. Overall, my view is this: JAX has many functions that are conceptually functions on Tensor1, but are defined for higher-order tensors, taking an axis parameter and having different default behaviors (sometimes the last axis, sometimes flattening the tensor).

An illustrative example for this is cross (not, yet in DimWit). Mathematically, an operation: vector x vector -> vector, so an operation on vectors. In JAX, we can provide two higher-dimensional tensors, and two (optional) axis arguments to say which vector(s) within these tensors, and even an output axis where to put the resulting vector(s). In DimWit the user "provides" these axis arguments with vmap or vapply, with cross being a function (Tensor1, Tensor1) => Tensor1, representing the mathematical scope.
This: When lifting JAX methods to DimWit, we should always be very skeptical about axis arguments. They are often a design flaw by JAX/numpy, that we can fix in DimWit.

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

marcelluethi added 4 commits September 5, 2026 06:19

add sort, cumsum, cumprod and diff methods

fb87b5e

add arc..., roundingops and isnan/isfinite

92fe377

add modulo operations

0e9e213

add nanToNum extension method

24b73aa

marcelluethi requested a review from benikm91 September 5, 2026 04:50

benikm91 reviewed Sep 6, 2026

View reviewed changes

cumsum/cumprod without label sum/factor over all elements

9344c83

benikm91 reviewed Sep 6, 2026

View reviewed changes

fix tests

e15a720

Uh oh!

Conversation

marcelluethi commented Sep 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

benikm91 Sep 6, 2026 •

edited

Loading