From 73bfc5812f045a16fc4297241923c47b06a01abd Mon Sep 17 00:00:00 2001 From: ppraneth Date: Mon, 9 Nov 2026 09:33:25 +0530 Subject: [PATCH] Add rewrite, using Theorem 1 for zero-preserving functions on sparse tensors --- src/function.rs | 6 +- src/index_functions.rs | 54 +++++++++++++++ src/lattice.rs | 6 ++ src/lib.rs | 5 +- src/rewrite.rs | 62 +++++++++++++++++ tests/index_functions.rs | 47 ++++++++++++- tests/rewrite.rs | 143 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 317 insertions(+), 6 deletions(-) create mode 100644 src/rewrite.rs create mode 100644 tests/rewrite.rs diff --git a/src/function.rs b/src/function.rs index baf5cb4..c546f88 100644 --- a/src/function.rs +++ b/src/function.rs @@ -43,8 +43,8 @@ impl Function for Rc { /// `f ∘ g`: applies `g`, then `f` to the result. pub struct Composition { - f: F, - g: G, + pub(crate) f: F, + pub(crate) g: G, } /// Composes `f` after `g`. Composing a tensor after an index map re-indexes it, @@ -89,7 +89,7 @@ where /// of its elements, repeats included, so it goes from the multisets of the /// domain of `f` to those of its codomain. pub struct MultiFunction { - f: F, + pub(crate) f: F, } pub fn multiset_lift(f: F) -> MultiFunction { diff --git a/src/index_functions.rs b/src/index_functions.rs index 273d5ce..2190c0e 100644 --- a/src/index_functions.rs +++ b/src/index_functions.rs @@ -331,10 +331,30 @@ impl IndexFunction for Where { } } +/// Whether a multi-index function sends each index to at most one index, which +/// makes a sum over it a sparse tensor (Definition 4 of the paper). The answer +/// may be a conservative `false`, never a wrong `true`. +pub trait Sparsity { + fn is_sparse(&self) -> bool; +} + +impl Sparsity for Rc { + fn is_sparse(&self) -> bool { + (**self).is_sparse() + } +} + pub struct EmptyMapping { signature: MultiIndexFunctionSignature, } +// Every index is sent to no index at all. +impl Sparsity for EmptyMapping { + fn is_sparse(&self) -> bool { + true + } +} + impl EmptyMapping { pub fn new(domain: HRect, codomain: HRect) -> Self { Self { @@ -374,6 +394,18 @@ impl COO { } } +// A partial coordinate matches every index sharing its prefix, so two stored +// coordinates land on a common index when one is a prefix of the other, equal +// ones included. Sorted, the coordinates starting with a given one come right +// after it, so only neighbours need comparing. +impl Sparsity for COO { + fn is_sparse(&self) -> bool { + let mut coords: Vec<&Vec> = self.inverse_coords.iter().collect(); + coords.sort(); + coords.windows(2).all(|pair| !pair[1].starts_with(pair[0])) + } +} + impl IndexFunction for COO { type Codomain = Multisets; @@ -434,6 +466,13 @@ impl IndexFunction for Diagonal { } } +// An index on the diagonal is sent to one index, and any other to none. +impl Sparsity for Diagonal { + fn is_sparse(&self) -> bool { + true + } +} + /// Projection dropping axis `dim`. /// e.g. with dim=1, f(m, n, p) = [m, p]^T pub struct Projection { @@ -516,6 +555,13 @@ impl IndexFunction for InvProjection { } } +// Each index is sent to the whole line along `dim`. +impl Sparsity for InvProjection { + fn is_sparse(&self) -> bool { + self.signature.codomain().set().shape()[self.dim] <= 1 + } +} + /// The affine map `x ↦ basis·x + offset`. `basis` has a row per dimension of the /// codomain and a column per dimension of the domain. Its entries and `offset` /// may be negative, as long as every index of the domain lands in the codomain. @@ -642,3 +688,11 @@ impl IndexFunction for InvAffine { &self.signature } } + +// Each index is sent to its preimage under the affine map, which holds at most +// one index when the map is one-to-one. +impl Sparsity for InvAffine { + fn is_sparse(&self) -> bool { + self.preimage.is_injective() + } +} diff --git a/src/lattice.rs b/src/lattice.rs index 1c7d4d3..d7e740c 100644 --- a/src/lattice.rs +++ b/src/lattice.rs @@ -38,6 +38,12 @@ impl Preimage { } } + /// Whether `x ↦ basis·x` is one-to-one, i.e. the basis has no integer + /// kernel. A kernel the box is too small to reach still counts. + pub(crate) fn is_injective(&self) -> bool { + self.h_pivots.len() == self.extents.len() + } + /// Every index `x` of the box with `basis·x = c`. pub(crate) fn solve(&self, c: &[i64]) -> Vec> { let mut z = vec![0; self.h_pivots.len()]; diff --git a/src/lib.rs b/src/lib.rs index 89ab9db..4c40008 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,7 @@ pub mod index_functions; mod lattice; pub mod multiset; pub mod reductions; +pub mod rewrite; pub mod sets; pub mod signature; pub mod tensor; @@ -21,13 +22,15 @@ pub use functions::{ }; pub use index_functions::{ Affine, COO, Concat, Diagonal, EmptyMapping, IndexFunction, InvAffine, InvProjection, - MultiIndexFunctionSignature, Permutation, Projection, Squeeze, Stride, Unsqueeze, Where, + MultiIndexFunctionSignature, Permutation, Projection, Sparsity, Squeeze, Stride, Unsqueeze, + Where, }; pub use multiset::Multiset; pub use reductions::{ Inf, InfReduction, Product, ProductReduction, Reduction, ReductionFunction, Sum, SumReduction, Sup, SupReduction, }; +pub use rewrite::{Rewrite, rewrite}; pub use sets::{HRect, IndexSet, JaggedDomain, Multisets, Set, Values}; pub use signature::{HRectSignature, Signature}; pub use tensor::{JaggedTensor, Physical, ReductionTensor, Tensor}; diff --git a/src/rewrite.rs b/src/rewrite.rs new file mode 100644 index 0000000..5588493 --- /dev/null +++ b/src/rewrite.rs @@ -0,0 +1,62 @@ +//! Rewriting a tensor into an equivalent one that is cheaper to evaluate, with +//! the identities between compositions of Section 4 of the paper. Only +//! Theorem 1 is used so far. + +use crate::function::{ + Composition, Function, MultiFunction, ValueFunction, compose, multiset_lift, +}; +use crate::index_functions::Sparsity; +use crate::reductions::{Reduction, Sum, SumReduction}; +use crate::sets::{Multisets, Values}; +use crate::tensor::ReductionTensor; + +/// A tensor that an identity of the paper can rewrite into an equivalent one. +pub trait Rewrite: Sized { + type Rewritten; + + /// The equivalent tensor, or `Err` with the tensor unchanged when the + /// identity does not apply to it. + fn rewrite(self) -> Result; +} + +/// Rewrites `tensor` into an equivalent one that is cheaper to evaluate, or +/// gives it back unchanged as `Err` when no identity applies to it. +pub fn rewrite(tensor: T) -> Result { + tensor.rewrite() +} + +/// Theorem 1: a function on values after a sparse tensor only has to be applied +/// to its stored values, `f ∘ Σ ∘ Ū ∘ ψ = Σ ∘ (f ∘ U)‾ ∘ ψ`. This holds when +/// `f(0) = 0` and `ψ` sends each index to at most one index, which is checked. +impl Rewrite for Composition, ReductionTensor, U, M>> +where + A: Clone + PartialEq, + SumReduction: Reduction, + U: Function>, + M: Function> + Sparsity, +{ + type Rewritten = ReductionTensor, Composition, U>, M>; + + fn rewrite(self) -> Result { + let zero = >::NEUTRAL; + if self.f.apply(&zero) != zero || !self.g.g.g.is_sparse() { + return Err(self); + } + let Composition { + f, + g: + Composition { + g: + Composition { + f: MultiFunction { f: u }, + g: psi, + }, + .. + }, + } = self; + Ok(compose( + Sum::new(), + compose(multiset_lift(compose(f, u)), psi), + )) + } +} diff --git a/tests/index_functions.rs b/tests/index_functions.rs index fa746a4..dfd93ca 100644 --- a/tests/index_functions.rs +++ b/tests/index_functions.rs @@ -2,8 +2,8 @@ use std::{assert_eq, collections::HashMap, rc::Rc}; use cten::{ Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexFunction, IndexSet, InvAffine, - InvProjection, JaggedDomain, Multiset, Permutation, Physical, Projection, Squeeze, Stride, - Unsqueeze, Where, + InvProjection, JaggedDomain, Multiset, Permutation, Physical, Projection, Sparsity, Squeeze, + Stride, Unsqueeze, Where, }; /// The indices of a multiset in sorted order, to compare multisets whose @@ -407,3 +407,46 @@ fn a_where_sends_each_index_to_the_part_its_condition_picks() { assert_eq!(picks.map(&[1]), vec![1, 1]); assert_eq!(picks.map(&[2]), vec![0, 2]); } + +// A partial coordinate matches every index sharing its prefix, so it meets every +// coordinate it is a prefix of. +#[test] +fn a_coo_is_sparse_when_no_two_coordinates_meet() { + let coo = |coords: Vec>| COO::new(HRect::new(vec![3, 3]), coords); + + assert!(coo(vec![vec![0, 2], vec![1, 0], vec![1, 2]]).is_sparse()); + assert!(!coo(vec![vec![1, 2], vec![0, 0], vec![1, 2]]).is_sparse()); + assert!(!coo(vec![vec![2], vec![2, 1]]).is_sparse()); + assert!(coo(vec![vec![2], vec![1, 1]]).is_sparse()); +} + +#[test] +fn diagonal_and_empty_index_functions_are_sparse() { + assert!(Diagonal::new(HRect::new(vec![3, 3])).is_sparse()); + assert!(EmptyMapping::new(HRect::new(vec![3]), HRect::new(vec![0])).is_sparse()); +} + +#[test] +fn an_inv_projection_is_sparse_only_along_an_axis_of_extent_one() { + assert!(InvProjection::new(HRect::new(vec![2, 1]), 1).is_sparse()); + assert!(!InvProjection::new(HRect::new(vec![2, 3]), 1).is_sparse()); +} + +#[test] +fn an_inv_affine_is_sparse_when_its_affine_map_is_one_to_one() { + let diagonal = InvAffine::new(Affine::new( + HRect::new(vec![3]), + vec![vec![1], vec![1]], + vec![0, 0], + HRect::new(vec![3, 3]), + )); + let projection = InvAffine::new(Affine::new( + HRect::new(vec![2, 3, 4]), + vec![vec![1, 0, 0], vec![0, 0, 1]], + vec![0, 0], + HRect::new(vec![2, 4]), + )); + + assert!(diagonal.is_sparse()); + assert!(!projection.is_sparse()); +} diff --git a/tests/rewrite.rs b/tests/rewrite.rs new file mode 100644 index 0000000..c3e3f1c --- /dev/null +++ b/tests/rewrite.rs @@ -0,0 +1,143 @@ +use std::cell::Cell; +use std::rc::Rc; + +use cten::{ + COO, HRect, Physical, ReductionTensor, Stride, Sum, Tensor, ValueFunction, compose, + diagonalize, rewrite, scale, sparse_coo_tensor, square, sum, +}; + +/// The sparse matrix of torch's `sparse_coo_tensor` documentation, which stores +/// three values: `[[0, 0, 3, 0], [4, 0, 5, 0]]`. +fn sparse() -> ReductionTensor { + sparse_coo_tensor( + vec![vec![0, 1, 1], vec![2, 0, 2]], + Physical::new(vec![3.0, 4.0, 5.0]), + vec![2, 4], + ) +} + +fn collect_2d(tensor: &dyn Tensor, rows: usize, cols: usize) -> Vec> { + (0..rows) + .map(|i| (0..cols).map(|j| tensor.apply(&[i, j])).collect()) + .collect() +} + +/// `x ↦ x²`, counting how many times it is applied. +fn counting_square(calls: &Rc>) -> ValueFunction { + let calls = calls.clone(); + ValueFunction::new(move |x: f32| { + calls.set(calls.get() + 1); + x * x + }) +} + +#[test] +fn theorem_1_rewrites_a_function_after_a_sparse_tensor_into_a_sparse_tensor() { + let Ok(rewritten) = rewrite(square(sparse())) else { + panic!("squaring keeps 0, so theorem 1 applies"); + }; + + assert_eq!( + collect_2d(&rewritten, 2, 4), + vec![vec![0.0, 0.0, 9.0, 0.0], vec![16.0, 0.0, 25.0, 0.0]] + ); + assert_eq!( + collect_2d(&rewritten, 2, 4), + collect_2d(&square(sparse()), 2, 4) + ); +} + +// The point of theorem 1: evaluated as written, the function is applied at every +// index of the matrix; rewritten, only at the three stored values. +#[test] +fn a_rewritten_tensor_applies_the_function_to_the_stored_values_only() { + let calls = Rc::new(Cell::new(0)); + collect_2d(&compose(counting_square(&calls), sparse()), 2, 4); + assert_eq!(calls.get(), 8); + + let calls = Rc::new(Cell::new(0)); + let Ok(rewritten) = rewrite(compose(counting_square(&calls), sparse())) else { + panic!("squaring keeps 0, so theorem 1 applies"); + }; + // Checking that the function keeps 0 applied it once. + calls.set(0); + collect_2d(&rewritten, 2, 4); + assert_eq!(calls.get(), 3); +} + +#[test] +fn a_rewritten_tensor_prints_like_the_original() { + let Ok(rewritten) = rewrite(square(sparse())) else { + panic!("squaring keeps 0, so theorem 1 applies"); + }; + + assert_eq!(rewritten.to_string(), square(sparse()).to_string()); +} + +#[test] +fn theorem_1_applies_to_other_functions_and_sparse_formats() { + let relu = || ValueFunction::new(|x: f32| x.max(0.0)); + let diagonal = || diagonalize(Physical::new(vec![-1.0, 2.0, -3.0]), 2); + let Ok(rewritten) = rewrite(compose(relu(), diagonal())) else { + panic!("relu keeps 0, so theorem 1 applies"); + }; + + assert_eq!( + collect_2d(&rewritten, 3, 3), + vec![ + vec![0.0, 0.0, 0.0], + vec![0.0, 2.0, 0.0], + vec![0.0, 0.0, 0.0] + ] + ); + assert_eq!( + collect_2d(&rewritten, 3, 3), + collect_2d(&compose(relu(), diagonal()), 3, 3) + ); +} + +// The result is a sparse tensor again, so a second function rewrites the same way. +#[test] +fn a_rewritten_tensor_can_be_rewritten_again() { + let Ok(squared) = rewrite(square(sparse())) else { + panic!("squaring keeps 0, so theorem 1 applies"); + }; + let Ok(rewritten) = rewrite(scale(squared, 2.0)) else { + panic!("scaling keeps 0, so theorem 1 applies"); + }; + + assert_eq!( + collect_2d(&rewritten, 2, 4), + vec![vec![0.0, 0.0, 18.0, 0.0], vec![32.0, 0.0, 50.0, 0.0]] + ); +} + +// Adding 1 does not keep 0: off the stored values, the tensor is 1 and not 0. The +// tensor comes back unchanged. +#[test] +fn theorem_1_needs_a_function_that_keeps_zero() { + let plus_one = ValueFunction::new(|x: f32| x + 1.0); + let Err(unchanged) = rewrite(compose(plus_one, sparse())) else { + panic!("adding 1 does not keep 0, so theorem 1 does not apply"); + }; + + assert_eq!( + collect_2d(&unchanged, 2, 4), + vec![vec![1.0, 1.0, 4.0, 1.0], vec![5.0, 1.0, 6.0, 1.0]] + ); +} + +// With several values at one index, squaring their sum is not summing their +// squares, so theorem 1 does not apply. +#[test] +fn theorem_1_needs_at_most_one_stored_value_per_index() { + let repeated = sparse_coo_tensor(vec![vec![1, 1]], Physical::new(vec![2.0, 3.0]), vec![3]); + let matrix = compose( + Physical::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + Stride::new(HRect::new(vec![2, 3]), vec![3, 1]), + ); + let row_sums = sum(matrix, 1); + + assert!(rewrite(square(repeated)).is_err()); + assert!(rewrite(square(row_sums)).is_err()); +}