Add commonly-used methods on tensors

core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala

@@ -79,9 +79,15 @@ object ShapeTypeHelpers:
   trait AxisReplacer[TensorShape <: Tuple, Axis, AxisReplacement] extends AxisInTensor[TensorShape, Axis]:
     type NewShape <: Tuple
 
-  object AxisReplacer:
+  object AxisReplacer extends AxisReplacerLowPriority:
     type Aux[S <: Tuple, A, AR, O <: Tuple] = AxisReplacer[S, A, AR] { type NewShape = O }
 
+    /** Replacing an axis by itself leaves the shape unchanged; needs only AxisIndex, so it also works for abstract shapes. */
+    given identity[S <: Tuple, A](using idx: AxisIndex[S, A]): AxisReplacer.Aux[S, A, A, S] = new AxisReplacer[S, A, A]:
+      def index: Int = idx.index
+      type NewShape = S
+
+  trait AxisReplacerLowPriority:
     given bridge[S <: Tuple, A, AR, O <: Tuple](using
         idx: AxisIndex[S, A],
         replacer: Replacer.Aux[S, A, AR, O]

core/src/main/scala/dimwit/tensor/Tensor.scala

@@ -10,6 +10,7 @@ import dimwit.tensor.Labels
 import dimwit.tensor.TensorOps.IsBoolean
 import dimwit.tensor.TensorOps.IsFloating
 import dimwit.tensor.TensorOps.IsInteger
+import dimwit.tensor.TensorOps.IsNumber
 import dimwit.tensor.TypedIndex
 import dimwit.tensor.VType
 import me.shadaj.scalapy.py
@@ -265,7 +266,8 @@ object Tensor0:
   def apply[V](jaxValue: Jax.PyDynamic): Tensor0[V] = Tensor(jaxValue)
 
 /** Companion object for Tensors of rank 1 (vectors).
-  *  Provides factory methods for creating tensors of rank 1 with various value types.
+  *  Provides factory methods for creating tensors of rank 1 with various value types,
+  *  and functions from a Tensor1 to a Tensor1 (e.g. `cumsum`, `sort`, `softmax`) that can be lifted to any shape with `vapply`.
   */
 object Tensor1:
 
@@ -295,6 +297,55 @@ object Tensor1:
   def apply[L: Label](axisExtent: AxisExtent[L]): Tensor.DefaultsFactory[Tuple1[L]] = Tensor.DefaultsFactory(Shape(axisExtent))
   def apply[L: Label, V](axisExtent: AxisExtent[L], vtype: VType[V]): Tensor.TypedFactory[Tuple1[L], V] = Tensor.TypedFactory(Shape(axisExtent), vtype)
 
+  // ---------------------------------------------------------
+  // Functions from a Tensor1 to a Tensor1.
+  // They can be lifted to any tensor shape with `vapply`, e.g. `t.vapply(Axis[A])(Tensor1.softmax)`.
+  // ---------------------------------------------------------
+
+  /** sorts the Tensor1 `t`. */
+  def sort[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnp.sort(t.jaxValue, axis = 0))
+
+  /** returns the indices that would sort the Tensor1 `t`. */
+  def argsort[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, Int32] =
+    Tensor(Jax.jnp.argsort(t.jaxValue, axis = 0))
+
+  /** computes the cumulative sum of the Tensor1 `t`. */
+  def cumsum[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnp.cumsum(t.jaxValue, axis = 0))
+
+  /** computes the cumulative product of the Tensor1 `t`. */
+  def cumprod[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnp.cumprod(t.jaxValue, axis = 0))
+
+  /** computes the cumulative maximum of the Tensor1 `t`. */
+  def cummax[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.lax.cummax(t.jaxValue, axis = 0))
+
+  /** computes the cumulative minimum of the Tensor1 `t`. */
+  def cummin[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.lax.cummin(t.jaxValue, axis = 0))
+
+  /** computes the cumulative log-sum-exp of the Tensor1 `t`, i.e. a numerically stable `log(cumsum(exp(t)))`. */
+  def logcumsumexp[L: Label, V: IsFloating](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.lax.cumlogsumexp(t.jaxValue, axis = 0))
+
+  /** computes the discrete difference of the Tensor1 `t`, reducing its size by one. */
+  def diff[L: Label, V: IsNumber](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnp.diff(t.jaxValue, axis = 0))
+
+  /** rolls the elements of the Tensor1 `t` by `shift` positions; elements shifted beyond the end re-appear at the start. */
+  def roll[L: Label, V](shift: Int)(t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnp.roll(t.jaxValue, shift = shift, axis = 0))
+
+  /** computes the softmax of the Tensor1 `t`. */
+  def softmax[L: Label, V: IsFloating](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnn.softmax(t.jaxValue, axis = 0))
+
+  /** computes the log of the softmax of the Tensor1 `t`, more stable than `softmax(t).log`. */
+  def logSoftmax[L: Label, V: IsFloating](t: Tensor1[L, V]): Tensor1[L, V] =
+    Tensor(Jax.jnn.log_softmax(t.jaxValue, axis = 0))
+
 /* Companion object for Tensors of rank 2 (matrices).
  *  Provides factory methods for creating tensors of rank 2 with various value types.
  */

core/src/main/scala/dimwit/tensor/TensorOps.scala

@@ -1,5 +1,6 @@
 package dimwit.tensor
 
+import dimwit.jax.Jax
 import dimwit.tensor.HasScalar
 import dimwit.tensor.Label
 import dimwit.tensor.Labels
@@ -33,6 +34,12 @@ object TensorOps:
   trait IsFloating[V] extends IsNumber[V], HasDType[V]:
     def dtype: DType
 
+    /** the largest finite value representable by V. */
+    def maxFinite: Tensor0[V] = Tensor0(Jax.jnp.array(Jax.jnp.finfo(dtype.jaxType).max, dtype = dtype.jaxType))
+
+    /** the smallest (most negative) finite value representable by V. */
+    def minFinite: Tensor0[V] = Tensor0(Jax.jnp.array(Jax.jnp.finfo(dtype.jaxType).min, dtype = dtype.jaxType))
+
   object IsFloating:
     def apply[V](using ev: IsFloating[V]): IsFloating[V] = ev
 
@@ -54,6 +61,7 @@ object TensorOps:
 
   export tensorops.ElementWiseOps.*
   export tensorops.ReductionOps.*
+  export tensorops.AlongAxisOps.*
   export tensorops.ContractionOps.*
   export tensorops.ConvolutionOps.*
   export tensorops.LinearAlgebraOps.*

core/src/main/scala/dimwit/tensor/tensorops/AlongAxisOps.scala

@@ -0,0 +1,72 @@
+package dimwit.tensor.tensorops
+
+import dimwit.jax.Jax
+import dimwit.tensor.Axis
+import dimwit.tensor.DType.Int32
+import dimwit.tensor.Label
+import dimwit.tensor.Labels
+import dimwit.tensor.ShapeTypeHelpers.AxisIndex
+import dimwit.tensor.ShapeTypeHelpers.AxisIndices
+import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes
+import dimwit.tensor.Tensor
+import dimwit.tensor.Tensor1
+import dimwit.tensor.TensorOps.IsFloating
+import dimwit.tensor.TensorOps.IsNumber
+import me.shadaj.scalapy.py.SeqConverters
+import dimwit.tensor.tensorops.FunctionalOps.vapply
+
+/** Operations along an axis that keep all axes of the tensor (unlike reductions, which remove them). */
+object AlongAxisOps:
+
+  extension [T <: Tuple: Labels, V](t: Tensor[T, V])
+
+    /** rolls the elements of `t` along the specified axis by `shift` positions. */
+    def roll[L: Label](axis: Axis[L], shift: Int)(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.roll(shift))
+
+  extension [T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])
+
+    /** Returns a tensor of indices that would sort `t` along the specified axes */
+    def argsort[Inputs <: Tuple](axes: Inputs)(using ev: AxisIndices[T, UnwrapAxes[Inputs]]): Tensor[T, Int32] = Tensor(Jax.jnp.argsort(t.jaxValue, axis = ev.indices.toPythonProxy))
+    def argsort[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, Int32] =
+      t.vapply(axis)(Tensor1.argsort)
+    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 AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.sort)
+    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 AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.cumsum)
+
+    /** computes the cumulative product of the tensor `t` along the specified axis. */
+    def cumprod[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.cumprod)
+
+    /** computes the cumulative maximum of the tensor `t` along the specified axis. */
+    def cummax[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.cummax)
+
+    /** computes the cumulative minimum of the tensor `t` along the specified axis. */
+    def cummin[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.cummin)
+
+    /** 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 AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.diff)
+
+  extension [T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V])
+
+    /** computes the cumulative log-sum-exp of the tensor `t` along the specified axis. */
+    def logcumsumexp[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.logcumsumexp)
+
+    /** computes the softmax of `t` along the specified axis. */
+    def softmax[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.softmax)
+
+    /** computes the log-softmax of `t` along the specified axis. */
+    def logSoftmax[L: Label](axis: Axis[L])(using AxisIndex[T, L]): Tensor[T, V] =
+      t.vapply(axis)(Tensor1.logSoftmax)

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])
@@ -173,6 +181,31 @@ 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))
+
+    /** 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
+        IsFloating[V]
+    )(
+        nan: Tensor0[V],
+        posInf: Tensor0[V] = IsFloating[V].maxFinite,
+        negInf: Tensor0[V] = IsFloating[V].minFinite
+    ): Tensor[T, V] =
+      Tensor(Jax.jnp.nan_to_num(t.jaxValue, nan = nan.jaxValue, posinf = posInf.jaxValue, neginf = negInf.jaxValue))
+
+    // 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))
 
     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] =

core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala

@@ -6,8 +6,6 @@ import dimwit.tensor.DType.Int32
 import dimwit.tensor.Label
 import dimwit.tensor.Labels
 import dimwit.tensor.ShapeTypeHelpers.AxesRemover
-import dimwit.tensor.ShapeTypeHelpers.AxisIndex
-import dimwit.tensor.ShapeTypeHelpers.AxisIndices
 import dimwit.tensor.ShapeTypeHelpers.AxisRemover
 import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes
 import dimwit.tensor.Tensor
@@ -47,11 +45,6 @@ object ReductionOps:
     def argmin[L: Label](axis: Axis[L])(using ev: AxisRemover[T, L], l: Labels[ev.RemainingAxes]): Tensor[ev.RemainingAxes, Int32] = Tensor(Jax.jnp.argmin(t.jaxValue, axis = ev.index))
     def argmin: Tensor0[Int32] = Tensor0(Jax.jnp.argmin(t.jaxValue))
 
-    /** Returns a tensor of indices that would sort `t` along the specified axes */
-    def argsort[Inputs <: Tuple](axes: Inputs)(using ev: AxisIndices[T, UnwrapAxes[Inputs]]): Tensor[T, Int32] = Tensor(Jax.jnp.argsort(t.jaxValue, axis = ev.indices.toPythonProxy))
-    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))
-
   // ---------------------------------------------------------
   // IsFloat operations (IsFloat or IsInt)
   // ---------------------------------------------------------

core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala

@@ -911,7 +911,3 @@ object StructuralOps:
         s"Cannot squeeze axis ${summon[Label[L]].name} of size ${tensor.shape.dimensions(ev.index)}"
       )
       Tensor(Jax.jnp.squeeze(tensor.jaxValue, axis = ev.index))
-
-  extension [L: Label, V](tensor: Tensor1[L, V])
-    def roll(shift: Int): Tensor1[L, V] =
-      Tensor(Jax.jnp.roll(tensor.jaxValue, shift = shift, axis = 0))