Add commonly-used methods on tensors

[marcelluethi]

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]

@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]

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]

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:

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

[marcelluethi]

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

[benikm91]

Comment out-of-date.

[benikm91]

Done.

[benikm91]

Remove default sort (last axis by convention).

[benikm91]

Done. I also removed argsort and I added tests that these methods do not exists (for Tensor2 objects) as we don't support a default axis if none is provided. We may support t1.sort (for Tensor1) in the future; tests are all written with Tensor2 for this case.

[benikm91]

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

[benikm91]

I tested it, it fails without it. Reason: The V: IsFloating is added after this parameter group, so IsFloating is not in scope for the default. So keep as is here.

[benikm91]

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.

[benikm91]

This is a bigger issue: In Tensor1 we have some functions, e.g., for softmax, but they are missing for ElementWiseOps and ReductionOps, etc. This is a inconsitency we should address. But maybe in a separate PR.

[benikm91]

See new comments.

Only the one def sort that defaults to axis=-1 must be removed, the other changes are optional, but consider them. If sort is removed I will approve.

[benikm91]

Looks good. I done the last few bits myself as discussed with @marcelluethi in our train ride 👍