diff --git a/core/src/main/scala/dimwit/nn/ActivationFunctions.scala b/core/src/main/scala/dimwit/nn/ActivationFunctions.scala
deleted file mode 100644
index ba344274..00000000
--- a/core/src/main/scala/dimwit/nn/ActivationFunctions.scala
+++ /dev/null
@@ -1,22 +0,0 @@
-package dimwit.nn
-
-import dimwit.jax.Jax
-import dimwit.python.PyBridge.liftPyTensor
-import dimwit.python.PyBridge.toPyTensor
-import dimwit.tensor.TensorOps.IsFloating
-import dimwit.tensor.*
-
-object ActivationFunctions:
-
-  def sigmoid[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] =
-    val x = Jax.jnn.sigmoid
-    liftPyTensor(Jax.jnn.sigmoid(toPyTensor(t)))
-
-  def relu[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] =
-    liftPyTensor(Jax.jnn.relu(toPyTensor(t)))
-
-  def gelu[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] =
-    liftPyTensor(Jax.jnn.gelu(toPyTensor(t)))
-
-  def softmax[L: Label, V: IsFloating](t: Tensor1[L, V]): Tensor1[L, V] =
-    liftPyTensor(Jax.jnn.softmax(toPyTensor(t), axis = 0))
diff --git a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala b/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala
index 37f54ebb..59aa934b 100644
--- a/core/src/main/scala/dimwit/tensor/ShapeTypeHelpers.scala
+++ b/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]
diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala
index 6ea683f9..ebb17ec5 100644
--- a/core/src/main/scala/dimwit/tensor/Tensor.scala
+++ b/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.
  */
diff --git a/core/src/main/scala/dimwit/tensor/TensorOps.scala b/core/src/main/scala/dimwit/tensor/TensorOps.scala
index ce1306c3..e05fa51c 100644
--- a/core/src/main/scala/dimwit/tensor/TensorOps.scala
+++ b/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.*
diff --git a/core/src/main/scala/dimwit/tensor/tensorops/AlongAxisOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/AlongAxisOps.scala
new file mode 100644
index 00000000..e7232dc5
--- /dev/null
+++ b/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)
diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala
index 37c98f33..c9d38697 100644
--- a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala
+++ b/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] =
diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala
index 7481004d..dc611104 100644
--- a/core/src/main/scala/dimwit/tensor/tensorops/ReductionOps.scala
+++ b/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)
   // ---------------------------------------------------------
diff --git a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala
index 5adbe2c2..bdc7b1c3 100644
--- a/core/src/main/scala/dimwit/tensor/tensorops/StructuralOps.scala
+++ b/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))
diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsAlongAxisSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsAlongAxisSuite.scala
new file mode 100644
index 00000000..970b7fc9
--- /dev/null
+++ b/core/src/test/scala/dimwit/tensor/TensorOpsAlongAxisSuite.scala
@@ -0,0 +1,174 @@
+package dimwit.tensor
+
+import dimwit.*
+
+class TensorOpsAlongAxisSuite extends DimwitTest:
+
+  val t2 = Tensor2(
+    Axis[A],
+    Axis[B]
+  ).fromArray(
+    Array(
+      Array(1.0f, 2.0f, 3.0f),
+      Array(4.0f, 5.0f, 6.0f)
+    )
+  )
+
+  val unsorted = Tensor2(
+    Axis[A],
+    Axis[B]
+  ).fromArray(
+    Array(
+      Array(1.0f, 3.0f, 2.0f),
+      Array(4.0f, 0.0f, 6.0f)
+    )
+  )
+
+  describe("Along Axis Ops"):
+    it("argsort"):
+      t2.argsort shouldEqual Tensor2(
+        Axis[A],
+        Axis[B]
+      ).fromArray(
+        Array(
+          Array(0, 1, 2),
+          Array(0, 1, 2)
+        )
+      )
+
+    it("argsort axis A"):
+      val res = t2.argsort(axis = Axis[A])
+      res shouldEqual Tensor2(
+        Axis[A],
+        Axis[B]
+      ).fromArray(
+        Array(
+          Array(0, 0, 0),
+          Array(1, 1, 1)
+        )
+      )
+
+    it("argsort axis B"):
+      val res = t2.argsort(axis = Axis[B])
+      res shouldEqual Tensor2(
+        Axis[A],
+        Axis[B]
+      ).fromArray(
+        Array(
+          Array(0, 1, 2),
+          Array(0, 1, 2)
+        )
+      )
+
+    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 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("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 axis A"):
+      val res = t2.diff(axis = Axis[A])
+      res shouldEqual Tensor.like(res).fromArray(Array(3.0f, 3.0f, 3.0f))
+
+    it("cummax axis B"):
+      val res = unsorted.cummax(axis = Axis[B])
+      res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 3.0f, 3.0f, 4.0f, 4.0f, 6.0f))
+
+    it("cummax axis A"):
+      val res = unsorted.cummax(axis = Axis[A])
+      res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 3.0f, 2.0f, 4.0f, 3.0f, 6.0f))
+
+    it("cummin axis B"):
+      val res = unsorted.cummin(axis = Axis[B])
+      res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 1.0f, 1.0f, 4.0f, 0.0f, 0.0f))
+
+    it("cummin axis A"):
+      val res = unsorted.cummin(axis = Axis[A])
+      res shouldEqual Tensor.like(res).fromArray(Array(1.0f, 3.0f, 2.0f, 1.0f, 0.0f, 2.0f))
+
+    it("logcumsumexp axis B"):
+      unsorted.logcumsumexp(axis = Axis[B]) should approxEqual(unsorted.exp.cumsum(Axis[B]).log, 1e-5f)
+
+    it("logcumsumexp axis A"):
+      unsorted.logcumsumexp(axis = Axis[A]) should approxEqual(unsorted.exp.cumsum(Axis[A]).log, 1e-5f)
+
+    it("Tensor1 functions lift with vapply"):
+      unsorted.vapply(Axis[B])(Tensor1.cummax) shouldEqual unsorted.cummax(Axis[B])
+
+    it("softmax axis B"):
+      val res = unsorted.softmax(axis = Axis[B])
+      res.sum(Axis[B]) should approxEqual(Tensor.like(res.sum(Axis[B])).fill(1.0f), 1e-5f)
+      res should approxEqual(unsorted.exp /! unsorted.exp.sum(Axis[B]), 1e-5f)
+
+    it("softmax axis A"):
+      val res = unsorted.softmax(axis = Axis[A])
+      res should approxEqual(unsorted.exp /! unsorted.exp.sum(Axis[A]), 1e-5f)
+
+    it("logSoftmax axis B"):
+      unsorted.logSoftmax(axis = Axis[B]) should approxEqual(unsorted.softmax(Axis[B]).log, 1e-5f)
+
+    it("roll with Tensor1.roll through vapply"):
+      unsorted.vapply(Axis[B])(Tensor1.roll(1)) shouldEqual unsorted.roll(Axis[B], shift = 1)
diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala
index 4e43664e..e69f05cc 100644
--- a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala
+++ b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala
@@ -31,6 +31,15 @@ class TensorOpsElementwiseSuite extends DimwitTest:
     it("abs"):
       t2.abs should approxEqual(Tensor.like(t2).fromArray(Array(1.0f, 0.0f, 1.0f, 4.0f)))
 
+    it("relu"):
+      t2.relu should approxEqual(Tensor.like(t2).fromArray(Array(0.0f, 0.0f, 1.0f, 4.0f)))
+
+    it("sigmoid"):
+      t2.sigmoid should approxEqual(Tensor.like(t2).fromArray(Array(0.26894143f, 0.5f, 0.7310586f, 0.98201376f)), tolerance = 1e-5f)
+
+    it("gelu"):
+      t2.gelu should approxEqual(Tensor.like(t2).fromArray(Array(-0.15880796f, 0.0f, 0.841192f, 3.9999299f)), tolerance = 1e-5f)
+
     it("sign"):
       t2.sign should approxEqual(Tensor.like(t2).fromArray(Array(-1.0f, 0.0f, 1.0f, 1.0f)))
 
@@ -53,6 +62,49 @@ 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(Tensor0(0f)) shouldEqual Tensor.like(t2).fromArray(Array(0.0f, Float.MaxValue, -Float.MaxValue, 1.0f))
+
+    it("maxFinite/minFinite of the floating type"):
+      IsFloating[Float32].maxFinite shouldEqual Tensor0(Float.MaxValue)
+      IsFloating[Float32].minFinite shouldEqual Tensor0(-Float.MaxValue)
+      IsFloating[Float64].maxFinite shouldEqual Tensor0(Double.MaxValue)
+
+    it("nanToNum with replacement values"):
+      val t = Tensor.like(t2).fromArray(Array(Float.NaN, Float.PositiveInfinity, Float.NegativeInfinity, 1.0f))
+      t.nanToNum(nan = Tensor0(-1.0f), posInf = Tensor0(100.0f), negInf = Tensor0(-100.0f)) shouldEqual
+        Tensor.like(t2).fromArray(Array(-1.0f, 100.0f, -100.0f, 1.0f))
+
+    it("nanToNum with only some replacement values"):
+      val t = Tensor.like(t2).fromArray(Array(Float.NaN, Float.PositiveInfinity, Float.NegativeInfinity, 1.0f))
+      t.nanToNum(nan = Tensor0(0f), posInf = Tensor0(100.0f)) shouldEqual Tensor.like(t2).fromArray(Array(0.0f, 100.0f, -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)))
 
@@ -75,6 +127,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))
 
diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala
index 9d4cdbd3..c65c83cb 100644
--- a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala
+++ b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala
@@ -133,41 +133,6 @@ class TensorOpsReductionSuite extends DimwitTest:
       val res = t2.argmin(axis = Axis[B])
       res shouldEqual Tensor.like(res).fromArray(Array(0, 0))
 
-    it("argsort"):
-      t2.argsort shouldEqual Tensor2(
-        Axis[A],
-        Axis[B]
-      ).fromArray(
-        Array(
-          Array(0, 1, 2),
-          Array(0, 1, 2)
-        )
-      )
-
-    it("argsort axis A"):
-      val res = t2.argsort(axis = Axis[A])
-      res shouldEqual Tensor2(
-        Axis[A],
-        Axis[B]
-      ).fromArray(
-        Array(
-          Array(0, 0, 0),
-          Array(1, 1, 1)
-        )
-      )
-
-    it("argsort axis B"):
-      val res = t2.argsort(axis = Axis[B])
-      res shouldEqual Tensor2(
-        Axis[A],
-        Axis[B]
-      ).fromArray(
-        Array(
-          Array(0, 1, 2),
-          Array(0, 1, 2)
-        )
-      )
-
   describe("Boolean Reductions"):
     it("all"):
       b2.all shouldEqual Tensor0(false)
diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala
index cd81dc3c..f9f0c4a3 100644
--- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala
+++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala
@@ -570,14 +570,14 @@ class TensorOpsStructureSuite extends DimwitTest:
     describe("roll"):
       it("tensor1"):
         val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f, 3.0f))
-        t1.roll(shift = 1) shouldEqual (Tensor1(Axis[A]).fromArray(Array(3.0f, 1.0f, 2.0f)))
-        t1.roll(shift = 2) shouldEqual (Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f, 1.0f)))
-        t1.roll(shift = 3) shouldEqual (t1)
+        t1.roll(Axis[A], shift = 1) shouldEqual (Tensor1(Axis[A]).fromArray(Array(3.0f, 1.0f, 2.0f)))
+        t1.roll(Axis[A], shift = 2) shouldEqual (Tensor1(Axis[A]).fromArray(Array(2.0f, 3.0f, 1.0f)))
+        t1.roll(Axis[A], shift = 3) shouldEqual (t1)
 
       it("tensor2"):
         val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f, 3.0f), Array(4.0f, 5.0f, 6.0f)))
-        t.vapply(Axis[B])(_.roll(shift = 1)) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 1.0f, 2.0f), Array(6.0f, 4.0f, 5.0f))))
-        t.vapply(Axis[A])(_.roll(shift = 1)) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(4.0f, 5.0f, 6.0f), Array(1.0f, 2.0f, 3.0f))))
+        t.roll(Axis[B], shift = 1) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 1.0f, 2.0f), Array(6.0f, 4.0f, 5.0f))))
+        t.roll(Axis[A], shift = 1) shouldEqual (Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(4.0f, 5.0f, 6.0f), Array(1.0f, 2.0f, 3.0f))))
 
   describe("chunk function"):
 
diff --git a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
index 827f881b..9f7b2bfb 100644
--- a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
+++ b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala
@@ -3,8 +3,6 @@ package dimwit.examples.basic
 import dimwit.Conversions.given
 import dimwit.*
 import dimwit.autodiff.*
-import dimwit.nn.ActivationFunctions.relu
-import dimwit.nn.ActivationFunctions.sigmoid
 import dimwit.optimizer.GradientDescent
 import dimwit.random.Random
 import dimwit.stats.Normal
@@ -24,7 +22,7 @@ object LogisticRegression:
       params.weights.dot(Axis[Feature])(input) + params.bias
 
     def probits(input: Tensor1[Feature, Float32]): Tensor0[Float32] =
-      sigmoid(logits(input))
+      logits(input).sigmoid
 
     def apply(input: Tensor1[Feature, Float32]): Tensor0[Bool] =
       logits(input) >= Tensor0(0f)
@@ -47,7 +45,7 @@ object LogisticRegression:
       val losses = zipvmap(Axis[Sample])(data, labels.asFloat32):
         case (sample, label) =>
           val logits = model.logits(sample)
-          relu(logits) - logits * label + ((-logits.abs).exp + 1f).log
+          logits.relu - logits * label + ((-logits.abs).exp + 1f).log
       losses.mean
 
   def main(args: Array[String]): Unit =
diff --git a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
index 39a7e8c1..faab4aca 100644
--- a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
+++ b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala
@@ -4,8 +4,6 @@ import dimwit.Conversions.given
 import dimwit.*
 import dimwit.tensortree.TreeOf.*
 import dimwit.autodiff.*
-import dimwit.nn.ActivationFunctions.relu
-import dimwit.nn.ActivationFunctions.sigmoid
 import dimwit.optimizer.GradientDescent
 import dimwit.python.PyBridge.toPyTensor
 import dimwit.random.Random
@@ -64,8 +62,8 @@ class Encoder(p: Encoder.Params):
   val logVarLayer = LinearLayer(p.logVarLayer)
 
   def apply(v: Tensor1[Pixel, Float32]): (Tensor1[Latent, Float32], Tensor1[Latent, Float32]) =
-    val h1 = relu(layer1(v))
-    val h2 = relu(layer2(h1))
+    val h1 = layer1(v).relu
+    val h2 = layer2(h1).relu
     val mean = meanLayer(h2)
     val logVar = logVarLayer(h2).clip(-10f, 10f)
     (mean, logVar)
@@ -85,9 +83,9 @@ class Decoder(p: Decoder.Params):
   val outputLayer = LinearLayer(p.outputLayer)
 
   def apply(v: Tensor1[Latent, Float32]): Tensor1[ReconstructedPixel, Float32] =
-    val h1 = relu(layer1(v))
-    val h2 = relu(layer2(h1))
-    sigmoid(outputLayer(h2))
+    val h1 = layer1(v).relu
+    val h2 = layer2(h1).relu
+    outputLayer(h2).sigmoid
 
 object Decoder:
   case class Params(