Make tensors generic on their value type

src/display.rs

@@ -3,7 +3,7 @@
 //! render the same way, by walking a shape.
 
 use std::fmt;
-use std::fmt::Write as _;
+use std::fmt::{Display, Write as _};
 
 use crate::domain::HRect;
 use crate::mapping::{BoxedMapping, BoxedPowersetMapping};
@@ -13,41 +13,43 @@ use crate::tensor::{
 
 /// A tensor whose domain is a hrect. Only these can be printed, since formatting
 /// walks a shape and visits every index of it.
-pub trait BoxedTensor: Tensor {
+pub trait BoxedTensor<A = f32>: Tensor<A> {
     fn hrect(&self) -> &HRect;
 }
 
 /// Renders a boxed tensor. Every `Display` implementation below defers to this,
 /// and each is written for the boxed instantiation of its tensor only.
-pub fn format_boxed(tensor: &dyn BoxedTensor) -> String {
+pub fn format_boxed<A: Display>(tensor: &dyn BoxedTensor<A>) -> String {
     format_tensor(tensor.hrect().shape(), &|index| tensor.get(index))
 }
 
-impl fmt::Display for Physical {
+impl<A: Clone + Display> Display for Physical<A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(&format_boxed(self))
     }
 }
 
-impl<M: BoxedMapping + ?Sized> fmt::Display for RightComposition<M> {
+impl<M: BoxedMapping + ?Sized, A: Display> Display for RightComposition<M, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(&format_boxed(self))
     }
 }
 
-impl<T: BoxedTensor + ?Sized> fmt::Display for LeftComposition<T> {
+impl<U: BoxedTensor<A> + ?Sized, A, B: Display> Display for LeftComposition<U, A, B> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(&format_boxed(self))
     }
 }
 
-impl<R: Reduction, M: BoxedPowersetMapping + ?Sized> fmt::Display for ReductionTensor<R, M> {
+impl<R: Reduction<A>, M: BoxedPowersetMapping + ?Sized, A: Display> Display
+    for ReductionTensor<R, M, A>
+{
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(&format_boxed(self))
     }
 }
 
-fn format_tensor(shape: &[usize], accessor: &dyn Fn(&[usize]) -> f32) -> String {
+fn format_tensor<A: Display>(shape: &[usize], accessor: &dyn Fn(&[usize]) -> A) -> String {
     if shape.is_empty() {
         return format!("tensor({})", accessor(&[]));
     }
@@ -70,11 +72,11 @@ fn format_tensor(shape: &[usize], accessor: &dyn Fn(&[usize]) -> f32) -> String
     out
 }
 
-fn max_width(
+fn max_width<A: Display>(
     shape: &[usize],
     index: &mut [usize],
     dim: usize,
-    accessor: &dyn Fn(&[usize]) -> f32,
+    accessor: &dyn Fn(&[usize]) -> A,
 ) -> usize {
     let mut width = 0;
     for i in 0..shape[dim] {
@@ -88,14 +90,14 @@ fn max_width(
     width
 }
 
-fn format_nested(
+fn format_nested<A: Display>(
     out: &mut String,
     shape: &[usize],
     index: &mut [usize],
     dim: usize,
     indent: &str,
     width: usize,
-    accessor: &dyn Fn(&[usize]) -> f32,
+    accessor: &dyn Fn(&[usize]) -> A,
 ) {
     let innermost = dim == shape.len() - 1;
     out.push('[');

src/functions.rs

@@ -3,6 +3,8 @@
 //! machinery of its own.
 
 use std::assert_eq;
+use std::fmt::Display;
+use std::ops::{Mul, Neg};
 use std::rc::Rc;
 
 use crate::display::BoxedTensor;
@@ -12,48 +14,65 @@ use crate::mapping::{
     Squeeze, Unsqueeze, Where,
 };
 use crate::tensor::{
-    JaggedTensor, LeftComposition, Physical, Product, Reduction, ReductionTensor, RightComposition,
-    Sum, Tensor,
+    JaggedTensor, LeftComposition, Physical, Product, ProductReduction, Reduction, ReductionTensor,
+    RightComposition, Sum, SumReduction, Tensor,
 };
 
-pub fn square<T: Tensor + ?Sized>(u: Rc<T>) -> LeftComposition<T> {
+pub fn square<A: Mul<Output = A> + Copy + 'static, U: Tensor<A> + ?Sized>(
+    u: Rc<U>,
+) -> LeftComposition<U, A> {
     LeftComposition::new(u, |x| x * x)
 }
 
-pub fn negate<T: Tensor + ?Sized>(u: Rc<T>) -> LeftComposition<T> {
+pub fn negate<A: Neg<Output = A> + 'static, U: Tensor<A> + ?Sized>(
+    u: Rc<U>,
+) -> LeftComposition<U, A> {
     LeftComposition::new(u, |x| -x)
 }
 
-pub fn scale<T: Tensor + ?Sized>(u: Rc<T>, factor: f32) -> LeftComposition<T> {
+pub fn scale<A: Mul<Output = A> + Copy + 'static, U: Tensor<A> + ?Sized>(
+    u: Rc<U>,
+    factor: A,
+) -> LeftComposition<U, A> {
     LeftComposition::new(u, move |x| x * factor)
 }
 
-pub fn pow<T: Tensor + ?Sized>(u: Rc<T>, exponent: f32) -> LeftComposition<T> {
-    LeftComposition::new(u, move |x| x.powf(exponent))
-}
-
-/// Checking the bounds here rather than leaving it to `f32::clamp` reports the
-/// mistake where it is made, not at the first read, and rules out the NaN
-/// bounds that would panic there too.
-pub fn clamp<T: Tensor + ?Sized>(u: Rc<T>, low: f32, high: f32) -> LeftComposition<T> {
+/// Holds values between `low` and `high`. Checking the bounds here reports the
+/// mistake where it is made, not at the first read, and rules out NaN bounds,
+/// which nothing would be below or above.
+pub fn clamp<A: PartialOrd + Copy + Display + 'static, U: Tensor<A> + ?Sized>(
+    u: Rc<U>,
+    low: A,
+    high: A,
+) -> LeftComposition<U, A> {
     assert!(low <= high, "clamp needs low <= high, got {low} and {high}");
-    LeftComposition::new(u, move |x| x.clamp(low, high))
+    LeftComposition::new(u, move |x| {
+        if x < low {
+            low
+        } else if x > high {
+            high
+        } else {
+            x
+        }
+    })
 }
 
 /// Reorders the axes of a tensor: `axes[i]` is the axis of `u` that becomes
 /// axis `i` of the result.
-pub fn permute<T: BoxedTensor + 'static>(
-    u: Rc<T>,
+pub fn permute<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
     axes: Vec<usize>,
-) -> RightComposition<Permutation> {
+) -> RightComposition<Permutation, A> {
     let codomain = u.hrect().clone();
     RightComposition::new(u, Rc::new(Permutation::new(codomain, axes)))
 }
 
 /// Swaps the last two axes and leaves any leading ones alone, so a stack of
 /// matrices transposes matrix by matrix. At rank two this is the plain
 /// transpose.
-pub fn transpose<T: BoxedTensor + 'static>(u: Rc<T>) -> RightComposition<Permutation> {
+pub fn transpose<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
+) -> RightComposition<Permutation, A> {
     let ndim = u.domain().ndim();
     assert!(ndim >= 2, "transpose needs at least two dimensions");
     let mut axes: Vec<usize> = (0..ndim).collect();
@@ -63,11 +82,11 @@ pub fn transpose<T: BoxedTensor + 'static>(u: Rc<T>) -> RightComposition<Permuta
 
 /// Moves the axis at `source` to `destination`; the axes left over keep their
 /// relative order.
-pub fn movedim<T: BoxedTensor + 'static>(
-    u: Rc<T>,
+pub fn movedim<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
     source: usize,
     destination: usize,
-) -> RightComposition<Permutation> {
+) -> RightComposition<Permutation, A> {
     let ndim = u.domain().ndim();
     assert!(source < ndim, "source axis {source} is out of range");
     assert!(
@@ -82,7 +101,10 @@ pub fn movedim<T: BoxedTensor + 'static>(
 
 /// Inserts an axis of extent 1 at `dim`, which can be anywhere from before the
 /// first axis to after the last.
-pub fn unsqueeze<T: BoxedTensor + 'static>(u: Rc<T>, dim: usize) -> RightComposition<Unsqueeze> {
+pub fn unsqueeze<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
+    dim: usize,
+) -> RightComposition<Unsqueeze, A> {
     let codomain = u.hrect().clone();
     RightComposition::new(u, Rc::new(Unsqueeze::new(codomain, dim)))
 }
@@ -127,7 +149,10 @@ impl Dims for &[usize] {
 /// Drops axes of extent 1, as torch's `squeeze` does: `squeeze(u, ())` drops all
 /// of them, `squeeze(u, 1)` only axis 1, and `squeeze(u, [0, 2])` axes 0 and 2.
 /// A named axis whose extent isn't 1 is left alone.
-pub fn squeeze<T: BoxedTensor + 'static>(u: Rc<T>, dims: impl Dims) -> RightComposition<Squeeze> {
+pub fn squeeze<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
+    dims: impl Dims,
+) -> RightComposition<Squeeze, A> {
     let codomain = u.hrect().clone();
     let ndim = codomain.shape().len();
     let named = dims.axes(ndim);
@@ -146,7 +171,10 @@ pub fn squeeze<T: BoxedTensor + 'static>(u: Rc<T>, dims: impl Dims) -> RightComp
 }
 
 /// Joins tensors end to end along `dim`; their other extents must agree.
-pub fn concatenate(tensors: Vec<Rc<dyn BoxedTensor>>, dim: usize) -> RightComposition<Concat> {
+pub fn concatenate<A: 'static>(
+    tensors: Vec<Rc<dyn BoxedTensor<A>>>,
+    dim: usize,
+) -> RightComposition<Concat, A> {
     assert!(!tensors.is_empty(), "concatenate needs at least one tensor");
     let underlying = JaggedTensor::new(tensors);
     let codomain = underlying.jagged_domain().clone();
@@ -155,7 +183,10 @@ pub fn concatenate(tensors: Vec<Rc<dyn BoxedTensor>>, dim: usize) -> RightCompos
 
 /// Joins tensors of one shape along a new axis at `dim`, by concatenating them
 /// once each has an axis of extent 1 there.
-pub fn stack(tensors: Vec<Rc<dyn BoxedTensor>>, dim: usize) -> RightComposition<Concat> {
+pub fn stack<A: 'static>(
+    tensors: Vec<Rc<dyn BoxedTensor<A>>>,
+    dim: usize,
+) -> RightComposition<Concat, A> {
     assert!(!tensors.is_empty(), "stack needs at least one tensor");
     let shape = tensors[0].hrect().clone();
     assert!(
@@ -166,20 +197,19 @@ pub fn stack(tensors: Vec<Rc<dyn BoxedTensor>>, dim: usize) -> RightComposition<
         .into_iter()
         .map(|tensor| {
             let unsqueeze = Unsqueeze::new(tensor.hrect().clone(), dim);
-            Rc::new(RightComposition::new(tensor, Rc::new(unsqueeze))) as Rc<dyn BoxedTensor>
+            Rc::new(RightComposition::new(tensor, Rc::new(unsqueeze))) as Rc<dyn BoxedTensor<A>>
         })
         .collect();
     concatenate(unsqueezed, dim)
 }
 
-/// Reads `input` where `condition` is not zero and `other` where it is. torch's
-/// `where` takes a boolean condition; cten's tensors hold floats, so zero stands
-/// for false. `where` is a Rust keyword, hence the underscore.
-pub fn where_(
-    condition: Rc<dyn BoxedTensor>,
-    input: Rc<dyn BoxedTensor>,
-    other: Rc<dyn BoxedTensor>,
-) -> RightComposition<Where> {
+/// Reads `input` where `condition` holds and `other` where it does not, as
+/// torch's `where` does. `where` is a Rust keyword, hence the underscore.
+pub fn where_<A: 'static>(
+    condition: Rc<dyn BoxedTensor<bool>>,
+    input: Rc<dyn BoxedTensor<A>>,
+    other: Rc<dyn BoxedTensor<A>>,
+) -> RightComposition<Where, A> {
     let underlying = JaggedTensor::new(vec![input, other]);
     let codomain = underlying.jagged_domain().clone();
     RightComposition::new(
@@ -191,7 +221,13 @@ pub fn where_(
 /// Spreads a vector over the diagonal of a tensor of the given rank: rank two
 /// gives a diagonal matrix, rank three a tensor whose only non-zero indices
 /// are of the form `[i, i, i]`.
-pub fn diagonalize<T: BoxedTensor + 'static>(u: Rc<T>, rank: usize) -> Sum<Diagonal> {
+pub fn diagonalize<A: 'static, U: BoxedTensor<A> + 'static>(
+    u: Rc<U>,
+    rank: usize,
+) -> Sum<Diagonal, A>
+where
+    SumReduction: Reduction<A>,
+{
     assert_eq!(u.domain().ndim(), 1, "diagonalize takes a vector");
     assert!(rank > 0, "a diagonal tensor has at least one dimension");
     let extent = u.hrect().shape()[0];
@@ -200,7 +236,7 @@ pub fn diagonalize<T: BoxedTensor + 'static>(u: Rc<T>, rank: usize) -> Sum<Diago
 
 /// Reducing over the empty set gives the reduction's neutral element at every
 /// index, so a constant tensor needs no underlying data at all.
-fn neutral<R: Reduction>(shape: Vec<usize>) -> ReductionTensor<R, EmptyMapping> {
+fn neutral<R: Reduction<f32>>(shape: Vec<usize>) -> ReductionTensor<R, EmptyMapping> {
     let mapping = EmptyMapping::new(HRect::new(shape), HRect::new(vec![0]));
     ReductionTensor::new(Rc::new(Physical::new(vec![])), Rc::new(mapping))
 }
@@ -231,11 +267,14 @@ pub fn eye(n: usize) -> Sum<Diagonal> {
 /// `indices`, and zero everywhere else. As in torch, `indices` has one row per
 /// dimension, so column `k` is the coordinate of `values[k]`, and values at a
 /// repeated coordinate add up.
-pub fn sparse_coo_tensor(
+pub fn sparse_coo_tensor<A: 'static>(
     indices: Vec<Vec<usize>>,
-    values: Rc<dyn BoxedTensor>,
+    values: Rc<dyn BoxedTensor<A>>,
     size: Vec<usize>,
-) -> Sum<COO> {
+) -> Sum<COO, A>
+where
+    SumReduction: Reduction<A>,
+{
     assert_eq!(
         indices.len(),
         size.len(),
@@ -273,30 +312,49 @@ pub fn sparse_coo_tensor(
 /// `physical` sits at `basis·x + offset`, and the tensor is zero off that
 /// lattice. `basis` has a row per dimension of `shape` and a column per
 /// dimension of `physical`. Values the lattice sends to one index add up.
-pub fn sparse_latticed_tensor<T: BoxedTensor + 'static>(
-    physical: Rc<T>,
+pub fn sparse_latticed_tensor<A: 'static, U: BoxedTensor<A> + 'static>(
+    physical: Rc<U>,
     basis: Vec<Vec<i64>>,
     offset: Vec<i64>,
     shape: Vec<usize>,
-) -> Sum<InvAffine> {
+) -> Sum<InvAffine, A>
+where
+    SumReduction: Reduction<A>,
+{
     let affine = Affine::new(physical.hrect().clone(), basis, offset, HRect::new(shape));
     Sum::new(physical, Rc::new(InvAffine::new(affine)))
 }
 
-pub fn sum<T: BoxedTensor + 'static>(input: Rc<T>, dim: usize) -> Sum<InvProjection> {
+pub fn sum<A: 'static, U: BoxedTensor<A> + 'static>(
+    input: Rc<U>,
+    dim: usize,
+) -> Sum<InvProjection, A>
+where
+    SumReduction: Reduction<A>,
+{
     let codomain = input.hrect().clone();
     Sum::new(input, Rc::new(InvProjection::new(codomain, dim)))
 }
 
-pub fn prod<T: BoxedTensor + 'static>(input: Rc<T>, dim: usize) -> Product<InvProjection> {
+pub fn prod<A: 'static, U: BoxedTensor<A> + 'static>(
+    input: Rc<U>,
+    dim: usize,
+) -> Product<InvProjection, A>
+where
+    ProductReduction: Reduction<A>,
+{
     let codomain = input.hrect().clone();
     Product::new(input, Rc::new(InvProjection::new(codomain, dim)))
 }
 
-pub fn mm<T1: BoxedTensor + 'static, T2: BoxedTensor + 'static>(
-    mat1: Rc<T1>,
-    mat2: Rc<T2>,
-) -> Sum<InvProjection> {
+pub fn mm<A: 'static, U: BoxedTensor<A> + 'static, V: BoxedTensor<A> + 'static>(
+    mat1: Rc<U>,
+    mat2: Rc<V>,
+) -> Sum<InvProjection, A>
+where
+    SumReduction: Reduction<A>,
+    ProductReduction: Reduction<A>,
+{
     assert_eq!(mat1.hrect().shape().len(), 2, "mat1 must be a matrix");
     assert_eq!(mat2.hrect().shape().len(), 2, "mat2 must be a matrix");
     assert_eq!(

src/lib.rs

@@ -12,8 +12,8 @@ pub mod tensor;
 pub use display::{BoxedTensor, format_boxed};
 pub use domain::{Domain, HRect, JaggedDomain};
 pub use functions::{
-    Dims, clamp, concatenate, diagonalize, eye, eye_rect, mm, movedim, negate, ones, permute, pow,
-    prod, scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, transpose,
+    Dims, clamp, concatenate, diagonalize, eye, eye_rect, mm, movedim, negate, ones, permute, prod,
+    scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, transpose,
     unsqueeze, where_, zeros,
 };
 pub use mapping::{