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: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

benikm91 · 2026-09-14T12:41:12Z

Maybe this approach is the best of both worlds (let's discuss tomorrow 👍 ):

Operations, which allow no or multiple axes in JAX, are define on Tensor like sum, mean, ... (as currently in main)

extension (t: Tensor[...])
  def sum(...)
t.sum
t.sum(Axis[A])

Operations, which require exactly one axis in JAX, are functions on vectors and should be defined on Tensor1.
However, we add an overload on Tensor that calls the function with vapply.
This achives both the right function scope and JAX-esk API:

// main branch
extension (t: Tensor1[...])
  def softmax: Tensor1[...] = ... // correct scoping

// not yet on main branch
extension (t: Tensor[...])
  def softmax(axis: Axis[A]): Tensor[...] = t.vapply(axis)(softmax) // syntax sugar
  
t.vapply(Axis[A])(softmax) // current
t.softmax(Axis[A]) // new option to do this

The implementation might be difficult due to Tensor1 being a Tensor, but should be possible with type classes. Let's first consider the API itself.

This consideration applys to sort, etc. from this PR, so let's decide here.

marcelluethi · 2026-09-15T06:02:11Z

So I understand that there are operations like diff, cumprod, softmax that conceptually are mappings from Tensor1 to Tensor1. These we could put in a object Tensor1Transform, e.g.:

object Tensor1Transform:
    def softmax[V : Floating](t : Tensor1[A, V]) : Tensor1[A, V] = ???  
    def diff[V](t : Tensor1[A, V]) : Tensor1[A, V] = ???

and in Tensor1Transform we have the convenience extension method:

extension t : Tensor[T, V : Floating]
    def softmax[L](axis : Axis[L]) : Tensor[T, V] = vapply(axis)(Tensor1Transform.softmax)

To use the methods, the user has two options:

Use t.softmax(Axis[L]) on an tensor of arbitrary dimension
Write vapply(Axis[L])(Tensor1Transform.softmax)

marcelluethi · 2026-09-19T15:21:58Z

@benikm91 I started doing the refactoring discussed above and incorporating your comments.

benikm91 · 2026-09-20T13:37:12Z

+    def isfinite: Tensor[T, Bool] = Tensor(Jax.jnp.isfinite(t.jaxValue))
+
+    /** replaces NaN by `nan`, +inf by `posInf` and -inf by `negInf`.
+      * By default, NaN becomes 0 and ±inf become the largest/smallest finite value of the dtype.


Comment out-of-date.

benikm91 · 2026-09-20T13:40:43Z

+    /** sorts the tensor `t` along the specified axis */
+    def sort[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.sort)
+    def sort: Tensor[T, V] = Tensor(Jax.jnp.sort(t.jaxValue))


Remove default sort (last axis by convention).

benikm91 · 2026-09-20T13:43:55Z

+    /** replaces NaN by `nan`, +inf by `posInf` and -inf by `negInf`.
+      * By default, NaN becomes 0 and ±inf become the largest/smallest finite value of the dtype.
+      */
+    def nanToNum(using


This IsFloating[V] can be removed. IsFloating evidence already in extention method.

benikm91 · 2026-09-20T13:45:59Z

+    // activation functions
+    def sigmoid: Tensor[T, V] = Tensor(Jax.jnn.sigmoid(t.jaxValue))
+    def relu: Tensor[T, V] = Tensor(Jax.jnn.relu(t.jaxValue))
+    def gelu: Tensor[T, V] = Tensor(Jax.jnn.gelu(t.jaxValue))


This changed syntax to t.relu. We need relu(t), as this is more natural for activation functions. I propose to put this in DimWit, here, but I am fine to move this also to DeepWit only.

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

Comment thread core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala

benikm91 reviewed Sep 6, 2026

View reviewed changes

Comment thread core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala

benikm91 reviewed Sep 6, 2026

View reviewed changes

marcelluethi added 2 commits September 6, 2026 17:25

fix tests

e15a720

Merge branch 'main' into more_elementary_jax_methods

713ba93

marcelluethi added 2 commits September 19, 2026 16:18

vapply based Tensor1/AlongAxisOps - move activation function

754f033

add arguments to nanToNum

1e07bda

benikm91 reviewed Sep 20, 2026

View reviewed changes

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!

Uh oh!

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!

benikm91 commented Sep 14, 2026

Uh oh!

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

Uh oh!

Uh oh!

marcelluethi commented Sep 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 Sep 20, 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!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

benikm91 Sep 6, 2026 •

edited

Loading

marcelluethi commented Sep 15, 2026 •

edited

Loading

benikm91 Sep 20, 2026 •

edited

Loading