From d3479bb557d365e60f9e83885e630e17fe65c1c0 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Thu, 5 Nov 2026 13:31:08 +0530 Subject: [PATCH 1/4] Add multiset-valued tensors and make reductions reduce multisets --- src/display.rs | 8 ++++- src/lib.rs | 11 ++++--- src/mapping.rs | 37 ++++++++++----------- src/multiset.rs | 74 ++++++++++++++++++++++++++++++++++++++++++ src/tensor.rs | 82 ++++++++++++++++++++++++++++++++++++++--------- tests/display.rs | 22 ++++++++++++- tests/logic.rs | 47 +++++++++++++++++++++++---- tests/mapping.rs | 30 +++++++++++------ tests/multiset.rs | 42 ++++++++++++++++++++++++ tests/tensor.rs | 48 +++++++++++++++++++++++++-- 10 files changed, 343 insertions(+), 58 deletions(-) create mode 100644 src/multiset.rs create mode 100644 tests/multiset.rs diff --git a/src/display.rs b/src/display.rs index de1bc02..1530a4d 100644 --- a/src/display.rs +++ b/src/display.rs @@ -8,7 +8,7 @@ use std::fmt::{Display, Write as _}; use crate::domain::HRect; use crate::mapping::{BoxedMapping, BoxedPowersetMapping}; use crate::tensor::{ - LeftComposition, Physical, Reduction, ReductionTensor, RightComposition, Tensor, + BarComposition, LeftComposition, Physical, Reduction, ReductionTensor, RightComposition, Tensor, }; /// A tensor whose domain is a hrect. Only these can be printed, since formatting @@ -41,6 +41,12 @@ impl + ?Sized, A, B: Display> Display for LeftComposition Display for BarComposition { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&format_boxed(self)) + } +} + impl, M: BoxedPowersetMapping + ?Sized, A: Display> Display for ReductionTensor { diff --git a/src/lib.rs b/src/lib.rs index 638af52..a88e64f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ pub mod domain; pub mod functions; mod lattice; pub mod mapping; +pub mod multiset; pub mod signature; pub mod tensor; @@ -18,11 +19,13 @@ pub use functions::{ }; pub use mapping::{ Affine, BoxedMapping, BoxedPowersetMapping, COO, Concat, Diagonal, EmptyMapping, IndexMapping, - InvAffine, InvProjection, Permutation, Powerset, PowersetMapping, Projection, Squeeze, Stride, - Unsqueeze, Where, + InvAffine, InvProjection, Permutation, PowersetMapping, Projection, Squeeze, Stride, Unsqueeze, + Where, }; +pub use multiset::Multiset; pub use signature::{HRectSignature, Signature}; pub use tensor::{ - Inf, InfReduction, JaggedTensor, LeftComposition, Physical, Product, ProductReduction, - Reduction, ReductionTensor, RightComposition, Sum, SumReduction, Sup, SupReduction, Tensor, + Bar, BarComposition, Inf, InfReduction, JaggedTensor, LeftComposition, Physical, Product, + ProductReduction, Reduction, ReductionTensor, RightComposition, Sum, SumReduction, Sup, + SupReduction, Tensor, }; diff --git a/src/mapping.rs b/src/mapping.rs index 61c0bcf..75e36f1 100644 --- a/src/mapping.rs +++ b/src/mapping.rs @@ -1,14 +1,11 @@ -use std::collections::HashSet; use std::rc::Rc; use crate::display::BoxedTensor; use crate::domain::{Domain, HRect, JaggedDomain}; use crate::lattice::Preimage; +use crate::multiset::Multiset; use crate::signature::{HRectSignature, Signature}; -/// A set of indices, as returned by a `PowersetMapping`. -pub type Powerset = HashSet>; - /// A mapping between two index sets, from its domain to its codomain. pub trait IndexMapping { fn map(&self, index: &[usize]) -> Vec; @@ -24,11 +21,11 @@ pub trait BoxedMapping: IndexMapping { fn hrect(&self) -> &HRect; } -/// A mapping sending each index of its domain to a set of indices of its -/// codomain. Note that the codomain is the set the returned indices are drawn -/// from, not its powerset. +/// A mapping sending each index of its domain to a multiset of indices of its +/// codomain, an index function in the paper's terms. Note that the codomain is +/// the set the returned indices are drawn from, not its multisets. pub trait PowersetMapping { - fn map(&self, index: &[usize]) -> Powerset; + fn map(&self, index: &[usize]) -> Multiset>; fn domain(&self) -> &dyn Domain; @@ -77,13 +74,13 @@ impl BoxedMapping for M { pub trait HRectPowersetMapping { type Codomain: Domain; - fn map(&self, index: &[usize]) -> Powerset; + fn map(&self, index: &[usize]) -> Multiset>; fn signature(&self) -> &Signature; } impl PowersetMapping for M { - fn map(&self, index: &[usize]) -> Powerset { + fn map(&self, index: &[usize]) -> Multiset> { HRectPowersetMapping::map(self, index) } @@ -408,8 +405,8 @@ impl EmptyMapping { impl HRectPowersetMapping for EmptyMapping { type Codomain = HRect; - fn map(&self, _index: &[usize]) -> Powerset { - Powerset::new() + fn map(&self, _index: &[usize]) -> Multiset> { + Multiset::default() } fn signature(&self) -> &HRectSignature { @@ -439,15 +436,15 @@ impl COO { impl HRectPowersetMapping for COO { type Codomain = HRect; - fn map(&self, index: &[usize]) -> Powerset { - let mut result = Powerset::new(); + fn map(&self, index: &[usize]) -> Multiset> { + let mut result = Multiset::default(); for (j, coords) in self.inverse_coords.iter().enumerate() { let matches = coords .iter() .enumerate() .all(|(i, coord)| *coord == index[i]); if matches { - result.insert(vec![j]); + result.push(vec![j]); } } result @@ -476,15 +473,15 @@ impl Diagonal { impl HRectPowersetMapping for Diagonal { type Codomain = HRect; - fn map(&self, index: &[usize]) -> Powerset { - let mut result = Powerset::new(); + fn map(&self, index: &[usize]) -> Multiset> { + let mut result = Multiset::default(); let first = index[0]; for coord in index { if *coord != first { return result; } } - result.insert(vec![first]); + result.push(vec![first]); result } @@ -557,7 +554,7 @@ impl InvProjection { impl HRectPowersetMapping for InvProjection { type Codomain = HRect; - fn map(&self, index: &[usize]) -> Powerset { + fn map(&self, index: &[usize]) -> Multiset> { (0..self.signature.codomain().shape()[self.dim]) .map(|coord| { let mut result = index.to_vec(); @@ -685,7 +682,7 @@ impl InvAffine { impl HRectPowersetMapping for InvAffine { type Codomain = HRect; - fn map(&self, index: &[usize]) -> Powerset { + fn map(&self, index: &[usize]) -> Multiset> { let target: Vec = index .iter() .zip(&self.offset) diff --git a/src/multiset.rs b/src/multiset.rs new file mode 100644 index 0000000..6f0956b --- /dev/null +++ b/src/multiset.rs @@ -0,0 +1,74 @@ +//! Multisets: collections in which an element may appear more than once, and +//! in which the multiplicity of each element counts but the order does not. + +use std::fmt; +use std::fmt::Display; + +/// An unordered collection of elements, repeats included. It is stored as a +/// vector, whose order carries no meaning. +#[derive(Clone, Debug)] +pub struct Multiset { + elements: Vec, +} + +impl Multiset { + pub fn new(elements: Vec) -> Self { + Self { elements } + } + + pub fn push(&mut self, element: A) { + self.elements.push(element); + } + + pub fn iter(&self) -> std::slice::Iter<'_, A> { + self.elements.iter() + } + + pub fn len(&self) -> usize { + self.elements.len() + } + + pub fn is_empty(&self) -> bool { + self.elements.is_empty() + } +} + +// Written out rather than derived, which would require `A: Default`. +impl Default for Multiset { + fn default() -> Self { + Self::new(Vec::new()) + } +} + +impl FromIterator for Multiset { + fn from_iter>(iter: I) -> Self { + Self::new(iter.into_iter().collect()) + } +} + +impl IntoIterator for Multiset { + type Item = A; + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.elements.into_iter() + } +} + +impl<'a, A> IntoIterator for &'a Multiset { + type Item = &'a A; + type IntoIter = std::slice::Iter<'a, A>; + + fn into_iter(self) -> Self::IntoIter { + self.elements.iter() + } +} + +/// Written as in the paper, with double braces: `{{1, 3, 3}}`. Padded as a whole, +/// so that multisets line up in a printed tensor as numbers do. +impl Display for Multiset { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let elements: Vec = self.elements.iter().map(A::to_string).collect(); + f.pad(&format!("{{{{{}}}}}", elements.join(", "))) + } +} diff --git a/src/tensor.rs b/src/tensor.rs index 72b63a0..6d83f42 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -4,6 +4,7 @@ use std::rc::Rc; use crate::display::BoxedTensor; use crate::domain::{Domain, HRect, JaggedDomain}; use crate::mapping::{BoxedMapping, BoxedPowersetMapping, IndexMapping, PowersetMapping}; +use crate::multiset::Multiset; /// A mapping from an index set to a value set. The index set is the tensor's /// domain, and the value set is its value type `A`, `f32` unless said @@ -161,6 +162,13 @@ impl + ?Sized, A, B> BoxedTensor for LeftComposition { const NEUTRAL: A; fn combine(a: A, b: A) -> A; + + /// Reduces a multiset to a single value: the neutral element for the empty + /// multiset, its element for a multiset of one, and otherwise its elements + /// combined in any order, as `combine` is associative and commutative. + fn reduce(multiset: Multiset) -> A { + multiset.into_iter().fold(Self::NEUTRAL, Self::combine) + } } pub struct SumReduction; @@ -208,21 +216,70 @@ reductions! { i64: 0, 1, i64::MIN, i64::MAX; } -// As in a right composition, the domain is the mapping's own domain, and the -// mapping type is a parameter so that reducing over a boxed mapping is known to -// give a boxed, printable tensor. -pub struct ReductionTensor { +/// The extension of a tensor `U` to multisets, written `Ū`: it maps a multiset of +/// indices to the multiset of the values of `U` at them, repeats included. Its +/// indices being multisets, it is not a tensor itself; composed with an index +/// function, it gives one. +pub struct Bar { u: Rc>, +} + +impl Bar { + pub fn new(u: Rc>) -> Self { + Self { u } + } + + pub fn get(&self, indices: &Multiset>) -> Multiset { + indices.iter().map(|index| self.u.get(index)).collect() + } +} + +// `Ū ∘ ψ`: at each index, the multiset of the values of `U` at the indices `ψ` +// maps it to. This is the multiset-valued tensor a reduction reduces. As in a +// right composition, the domain is the mapping's own domain, and the mapping +// type is a parameter so that a boxed mapping is known to give a boxed, +// printable tensor. +pub struct BarComposition { + bar: Bar, mapping: Rc, - _marker: PhantomData R>, } -impl, M: PowersetMapping + ?Sized, A> ReductionTensor { +impl BarComposition { pub fn new(u: Rc>, mapping: Rc) -> Self { check_composable(mapping.codomain(), &*u); Self { - u, + bar: Bar::new(u), mapping, + } + } +} + +impl Tensor> for BarComposition { + fn get(&self, index: &[usize]) -> Multiset { + self.bar.get(&self.mapping.map(index)) + } + + fn domain(&self) -> &dyn Domain { + self.mapping.domain() + } +} + +impl BoxedTensor> for BarComposition { + fn hrect(&self) -> &HRect { + self.mapping.hrect() + } +} + +// `Ω ∘ Ū ∘ ψ`: the reduction `R` of the multisets that a bar composition gathers. +pub struct ReductionTensor { + gathered: BarComposition, + _marker: PhantomData R>, +} + +impl, M: PowersetMapping + ?Sized, A> ReductionTensor { + pub fn new(u: Rc>, mapping: Rc) -> Self { + Self { + gathered: BarComposition::new(u, mapping), _marker: PhantomData, } } @@ -230,16 +287,11 @@ impl, M: PowersetMapping + ?Sized, A> ReductionTensor { impl, M: PowersetMapping + ?Sized, A> Tensor for ReductionTensor { fn get(&self, index: &[usize]) -> A { - let indices = self.mapping.map(index); - let mut result = R::NEUTRAL; - for i in &indices { - result = R::combine(result, self.u.get(i)); - } - result + R::reduce(self.gathered.get(index)) } fn domain(&self) -> &dyn Domain { - self.mapping.domain() + self.gathered.domain() } } @@ -247,7 +299,7 @@ impl, M: BoxedPowersetMapping + ?Sized, A> BoxedTensor for ReductionTensor { fn hrect(&self) -> &HRect { - self.mapping.hrect() + self.gathered.hrect() } } diff --git a/tests/display.rs b/tests/display.rs index e81aba4..b28fe8d 100644 --- a/tests/display.rs +++ b/tests/display.rs @@ -1,7 +1,8 @@ use std::rc::Rc; use cten::{ - COO, Diagonal, HRect, Physical, RightComposition, Stride, Sum, Sup, Tensor, format_boxed, + BarComposition, COO, Diagonal, HRect, Physical, RightComposition, Stride, Sum, Sup, Tensor, + format_boxed, }; fn physical() -> Rc { @@ -22,6 +23,25 @@ fn tensors_of_other_value_types_print_the_same_way() { assert_eq!(flags.to_string(), "tensor([ true, false])"); } +// The diagonal matrix of the paper's sparse matrices figure, before and after +// its multisets of values are summed. +#[test] +fn a_multiset_valued_tensor_prints_its_multisets() { + let stored = Rc::new(Physical::new(vec![1.0, 2.0, 4.0])); + let diagonal = Rc::new(Diagonal::new(HRect::new(vec![3, 3]))); + let before = BarComposition::new(stored.clone(), diagonal.clone()); + let after = Sum::new(stored, diagonal); + + assert_eq!( + before.to_string(), + "tensor([[{{1}}, {{}}, {{}}],\n [ {{}}, {{2}}, {{}}],\n [ {{}}, {{}}, {{4}}]])" + ); + assert_eq!( + after.to_string(), + "tensor([[1, 0, 0],\n [0, 2, 0],\n [0, 0, 4]])" + ); +} + #[test] fn a_rank_two_tensor_prints_one_row_per_line_aligned_under_the_bracket() { let strided = RightComposition::new( diff --git a/tests/logic.rs b/tests/logic.rs index 9557ae5..d661468 100644 --- a/tests/logic.rs +++ b/tests/logic.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - Affine, COO, Diagonal, EmptyMapping, HRect, Inf, LeftComposition, Physical, PowersetMapping, - Product, RightComposition, Stride, Sum, Sup, Tensor, + Affine, COO, Diagonal, EmptyMapping, HRect, HRectSignature, Inf, LeftComposition, Multiset, + Physical, PowersetMapping, Product, RightComposition, Stride, Sum, Sup, Tensor, }; fn physical() -> Rc { @@ -162,9 +162,9 @@ fn each_reduction_falls_back_to_its_neutral_element() { assert_eq!(inf.get(&off_support), f32::INFINITY); } -// The demo's COO patterns all use distinct coordinates, so every Powerset it -// builds holds at most one index and the fold in ReductionTensor never runs. -// These patterns deliberately store several positions at one coordinate. +// The demo's COO patterns all use distinct coordinates, so every multiset it +// builds holds at most one index and a reduction has nothing to combine. These +// patterns deliberately store several positions at one coordinate. #[test] fn reductions_accumulate_over_every_matching_position() { let overlapping = || { @@ -284,7 +284,7 @@ fn a_powerset_mapping_can_be_used_directly() { let selected = diagonal.map(&[4, 4]); assert_eq!(selected.len(), 1); - assert!(selected.contains([4].as_slice())); + assert!(selected.iter().any(|index| index == &[4])); assert!(diagonal.map(&[4, 5]).is_empty()); } @@ -310,3 +310,38 @@ fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { let read: Vec = (0..3).map(|i| diagonal.get(&[i])).collect(); assert_eq!(read, vec![1.0, 5.0, 9.0]); } + +/// An index function sending each index to the multiset holding it twice. +struct Twice { + signature: HRectSignature, +} + +// Named in full, as importing it would make `map` ambiguous in the tests above. +impl cten::mapping::HRectPowersetMapping for Twice { + type Codomain = HRect; + + fn map(&self, index: &[usize]) -> Multiset> { + Multiset::new(vec![index.to_vec(), index.to_vec()]) + } + + fn signature(&self) -> &HRectSignature { + &self.signature + } +} + +// Index functions map to multisets, so an index can be met more than once, and +// a reduction then counts its value as many times. +#[test] +fn a_reduction_counts_an_index_as_often_as_it_is_repeated() { + let twice = Rc::new(Twice { + signature: HRectSignature::new(HRect::new(vec![6]), HRect::new(vec![6])), + }); + let doubled = Sum::new(physical(), twice.clone()); + let squared = Product::new(physical(), twice); + + for i in 0..6 { + let value = (i + 1) as f32; + assert_eq!(doubled.get(&[i]), 2.0 * value); + assert_eq!(squared.get(&[i]), value * value); + } +} diff --git a/tests/mapping.rs b/tests/mapping.rs index 8e2888d..ab81889 100644 --- a/tests/mapping.rs +++ b/tests/mapping.rs @@ -2,10 +2,18 @@ use std::{assert_eq, collections::HashMap, rc::Rc}; use cten::{ Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexMapping, InvAffine, InvProjection, - JaggedDomain, Permutation, Physical, Powerset, PowersetMapping, Projection, Squeeze, Stride, + JaggedDomain, Multiset, Permutation, Physical, PowersetMapping, Projection, Squeeze, Stride, Unsqueeze, Where, }; +/// The indices of a multiset in sorted order, to compare multisets whose +/// elements may come out in different orders. +fn sorted(indices: Multiset>) -> Vec> { + let mut elements: Vec> = indices.into_iter().collect(); + elements.sort(); + elements +} + #[test] fn a_stride_goes_from_its_hrect_to_the_flat_positions_it_reaches() { let stride = Stride::new(HRect::new(vec![3, 2]), vec![1, 3]); @@ -66,8 +74,8 @@ fn an_inv_projection_sends_each_index_to_the_line_along_its_dim() { assert!(!inv_proj.domain().contains(&[1, 4])); assert!(inv_proj.codomain().contains(&[1, 2, 3])); assert_eq!( - inv_proj.map(&[1, 3]), - Powerset::from([vec![1, 0, 3], vec![1, 1, 3], vec![1, 2, 3]]) + sorted(inv_proj.map(&[1, 3])), + vec![vec![1, 0, 3], vec![1, 1, 3], vec![1, 2, 3]] ); } @@ -127,8 +135,8 @@ fn an_inv_affine_sends_each_index_to_the_lattice_points_on_it() { assert!(inv_affine.domain().contains(&[8])); assert!(!inv_affine.domain().contains(&[9])); assert!(inv_affine.codomain().contains(&[2])); - assert_eq!(inv_affine.map(&[4]), Powerset::from([vec![1]])); - assert_eq!(inv_affine.map(&[5]), Powerset::new()); + assert_eq!(sorted(inv_affine.map(&[4])), vec![vec![1]]); + assert!(inv_affine.map(&[5]).is_empty()); } // Dropping the middle axis is affine too, and sends several indices onto each @@ -145,7 +153,10 @@ fn an_inv_affine_of_a_projection_matches_the_inv_projection() { for m in 0..2 { for p in 0..4 { - assert_eq!(inv_affine.map(&[m, p]), inv_proj.map(&[m, p])); + assert_eq!( + sorted(inv_affine.map(&[m, p])), + sorted(inv_proj.map(&[m, p])) + ); } } } @@ -174,7 +185,8 @@ fn assert_inv_affine_matches_a_scan( physical: Vec, shape: Vec, ) { - let mut expected: HashMap, Powerset> = HashMap::new(); + // Scanned in row-major order, so each list of indices comes out sorted. + let mut expected: HashMap, Vec>> = HashMap::new(); for x in indices(&physical) { let y = basis .iter() @@ -188,7 +200,7 @@ fn assert_inv_affine_matches_a_scan( (start + coord) as usize }) .collect(); - expected.entry(y).or_default().insert(x); + expected.entry(y).or_default().push(x); } let inv_affine = InvAffine::new(Affine::new( HRect::new(physical), @@ -198,7 +210,7 @@ fn assert_inv_affine_matches_a_scan( )); for y in indices(&shape) { let lands = expected.remove(&y).unwrap_or_default(); - assert_eq!(inv_affine.map(&y), lands, "at {y:?}"); + assert_eq!(sorted(inv_affine.map(&y)), lands, "at {y:?}"); } } diff --git a/tests/multiset.rs b/tests/multiset.rs new file mode 100644 index 0000000..b96dd42 --- /dev/null +++ b/tests/multiset.rs @@ -0,0 +1,42 @@ +use cten::{InfReduction, Multiset, ProductReduction, Reduction, SumReduction, SupReduction}; + +#[test] +fn a_multiset_prints_with_double_braces() { + assert_eq!(Multiset::new(vec![1, 3, 3]).to_string(), "{{1, 3, 3}}"); + assert_eq!(Multiset::::default().to_string(), "{{}}"); +} + +#[test] +fn a_multiset_keeps_every_repeat() { + let multiset: Multiset = [3, 1, 3].into_iter().collect(); + + assert_eq!(multiset.len(), 3); + assert!(!multiset.is_empty()); + assert!(Multiset::::default().is_empty()); +} + +// The example of issue #18: the sum counts the 3 twice. +#[test] +fn a_reduction_reduces_a_multiset_to_one_value() { + assert_eq!( + SumReduction::reduce(Multiset::new(vec![1.0, 3.0, 3.0])), + 7.0 + ); + assert_eq!(SumReduction::reduce(Multiset::new(vec![1, 3, 3])), 7); + assert_eq!(ProductReduction::reduce(Multiset::new(vec![2, 3, 3])), 18); + assert_eq!(SupReduction::reduce(Multiset::new(vec![2, 7, 5])), 7); + assert_eq!(InfReduction::reduce(Multiset::new(vec![2, 7, 5])), 2); +} + +// Reducing the empty multiset gives the neutral element, and reducing a single +// element gives that element. +#[test] +fn a_reduction_of_no_value_or_one_value_is_as_expected() { + assert_eq!(SumReduction::reduce(Multiset::::default()), 0.0); + assert_eq!(ProductReduction::reduce(Multiset::::default()), 1.0); + assert_eq!( + SupReduction::reduce(Multiset::::default()), + f32::NEG_INFINITY + ); + assert_eq!(InfReduction::reduce(Multiset::new(vec![4.0])), 4.0); +} diff --git a/tests/tensor.rs b/tests/tensor.rs index aa376e7..56de005 100644 --- a/tests/tensor.rs +++ b/tests/tensor.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - BoxedTensor, COO, Diagonal, HRect, JaggedTensor, LeftComposition, Physical, RightComposition, - Stride, Sum, Tensor, + Bar, BarComposition, BoxedTensor, COO, Diagonal, HRect, InvProjection, JaggedTensor, + LeftComposition, Multiset, Physical, RightComposition, Stride, Sum, Tensor, }; fn physical() -> Rc { @@ -68,6 +68,50 @@ fn a_left_composition_takes_its_domain_from_the_tensor_it_transforms() { assert_eq!(squared.get(&[2, 1]), 36.0); } +// A multiset of indices reads a multiset of values, repeats included. +#[test] +fn a_bar_reads_a_multiset_of_values_at_a_multiset_of_indices() { + let bar = Bar::new(physical()); + let read = bar.get(&Multiset::new(vec![vec![0], vec![2], vec![2]])); + + assert_eq!(read.into_iter().collect::>(), vec![1.0, 3.0, 3.0]); + assert!(bar.get(&Multiset::default()).is_empty()); +} + +/// The matrix of the paper's row sums example, `[[1, 2, 3], [0, 4, 1], [2, 2, 5]]`. +fn figure_matrix() -> Rc> { + Rc::new(RightComposition::new( + Rc::new(Physical::new(vec![ + 1.0, 2.0, 3.0, 0.0, 4.0, 1.0, 2.0, 2.0, 5.0, + ])), + Rc::new(Stride::new(HRect::new(vec![3, 3]), vec![3, 1])), + )) +} + +// Before reduction, each row index reads the multiset of its row; reducing those +// multisets with a sum gives the row sums. +#[test] +fn a_bar_composition_is_a_reduction_before_the_reduction() { + let rows = Rc::new(InvProjection::new(HRect::new(vec![3, 3]), 1)); + let gathered = BarComposition::new(figure_matrix(), rows.clone()); + let row_sums = Sum::new(figure_matrix(), rows); + + assert_eq!(gathered.hrect().shape(), &[3]); + let before: Vec> = (0..3) + .map(|i| gathered.get(&[i]).into_iter().collect()) + .collect(); + assert_eq!( + before, + vec![ + vec![1.0, 2.0, 3.0], + vec![0.0, 4.0, 1.0], + vec![2.0, 2.0, 5.0] + ] + ); + let after: Vec = (0..3).map(|i| row_sums.get(&[i])).collect(); + assert_eq!(after, vec![6.0, 5.0, 9.0]); +} + #[test] fn a_physical_tensor_holds_values_of_any_type() { let ints = Physical::new(vec![3i64, -1]); From cd47a7f2e67863998e81cfd44b1128ed6cc0a093 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Thu, 5 Nov 2026 15:40:22 +0530 Subject: [PATCH 2/4] Rebuild the core on a single Function trait with compose and multiset_lift --- src/display.rs | 44 +++------ src/function.rs | 178 +++++++++++++++++++++++++++++++++ src/functions.rs | 218 ++++++++++++++++++++++------------------- src/lib.rs | 14 +-- src/main.rs | 22 +++-- src/mapping.rs | 138 +++++++++----------------- src/tensor.rs | 239 +++++++++++---------------------------------- tests/display.rs | 46 +++++---- tests/functions.rs | 167 ++++++++++++++++--------------- tests/logic.rs | 168 ++++++++++++++++++------------- tests/mapping.rs | 110 ++++++++++----------- tests/multiset.rs | 23 ++--- tests/tensor.rs | 152 ++++++++++++++++------------ 13 files changed, 793 insertions(+), 726 deletions(-) create mode 100644 src/function.rs diff --git a/src/display.rs b/src/display.rs index 1530a4d..1d29377 100644 --- a/src/display.rs +++ b/src/display.rs @@ -2,25 +2,17 @@ //! implementations live here rather than next to each tensor, since they all //! render the same way, by walking a shape. +use std::borrow::Borrow; use std::fmt; use std::fmt::{Display, Write as _}; -use crate::domain::HRect; -use crate::mapping::{BoxedMapping, BoxedPowersetMapping}; -use crate::tensor::{ - BarComposition, LeftComposition, Physical, Reduction, ReductionTensor, RightComposition, 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 { - fn hrect(&self) -> &HRect; -} +use crate::function::{Boxed, Compose, Function}; +use crate::tensor::{BoxedTensor, Physical}; /// Renders a boxed tensor. Every `Display` implementation below defers to this, -/// and each is written for the boxed instantiation of its tensor only. +/// and each is written for boxed tensors only. pub fn format_boxed(tensor: &dyn BoxedTensor) -> String { - format_tensor(tensor.hrect().shape(), &|index| tensor.get(index)) + format_tensor(tensor.hrect().shape(), &|index| tensor.apply(index)) } impl Display for Physical { @@ -29,26 +21,12 @@ impl Display for Physical { } } -impl Display for RightComposition { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&format_boxed(self)) - } -} - -impl + ?Sized, A, B: Display> Display for LeftComposition { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&format_boxed(self)) - } -} - -impl Display for BarComposition { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&format_boxed(self)) - } -} - -impl, M: BoxedPowersetMapping + ?Sized, A: Display> Display - for ReductionTensor +// A composition is a boxed tensor when the function applied first is one on a +// hrect and the one applied last gives values that print. +impl> Display for Compose +where + G::Output: Borrow, + F::Output: Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(&format_boxed(self)) diff --git a/src/function.rs b/src/function.rs new file mode 100644 index 0000000..eb19662 --- /dev/null +++ b/src/function.rs @@ -0,0 +1,178 @@ +//! The building bricks: functions, their composition, and their extension to +//! multisets. Tensors, index maps, index functions and reductions are all +//! functions, and every operation on them is built from these. + +use std::borrow::Borrow; +use std::rc::Rc; + +use crate::domain::{Domain, HRect}; +use crate::multiset::Multiset; + +/// A function from `Input` to `Output`. A tensor is a function from indices to +/// values, an index map one from indices to indices, an index function one from +/// indices to multisets of indices, and a reduction one from multisets of values +/// to values. +pub trait Function { + type Input: ?Sized; + type Output; + + fn apply(&self, x: &Self::Input) -> Self::Output; + + /// For a function on indices, or on multisets of indices, the set those + /// indices are drawn from. A function on values has none. + fn domain(&self) -> Option<&dyn Domain> { + None + } + + /// For a function to indices, or to multisets of indices, the set those + /// indices are drawn from. A function to values has none. + fn codomain(&self) -> Option<&dyn Domain> { + None + } +} + +/// A function whose domain is a hrect. Only these can be printed, since +/// formatting walks a shape and visits every index of it. +pub trait Boxed: Function { + fn hrect(&self) -> &HRect; +} + +// A shared function is the same function, so that one can be used in several +// compositions, or have its type erased as an `Rc>`. +impl Function for Rc { + type Input = T::Input; + type Output = T::Output; + + fn apply(&self, x: &Self::Input) -> Self::Output { + (**self).apply(x) + } + + fn domain(&self) -> Option<&dyn Domain> { + (**self).domain() + } + + fn codomain(&self) -> Option<&dyn Domain> { + (**self).codomain() + } +} + +impl Boxed for Rc { + fn hrect(&self) -> &HRect { + (**self).hrect() + } +} + +/// `f ∘ g`: applies `g`, then `f` to the result. +pub struct Compose { + f: F, + g: G, +} + +/// Composes `f` after `g`. Composing a tensor after an index map re-indexes it, +/// and a function on values after a tensor applies it element-wise. Where both +/// sides are on indices, the indices `g` lands on must be those `f` takes, which +/// is checked here as far as it can be without enumerating a domain. +pub fn compose(f: F, g: G) -> Compose +where + G::Output: Borrow, +{ + if let (Some(codomain), Some(domain)) = (g.codomain(), f.domain()) { + assert_eq!( + codomain.ndim(), + domain.ndim(), + "the inner function lands on {}-dimensional indices, but the outer one is indexed by {} coordinates", + codomain.ndim(), + domain.ndim() + ); + } + Compose { f, g } +} + +impl Function for Compose +where + G::Output: Borrow, +{ + type Input = G::Input; + type Output = F::Output; + + fn apply(&self, x: &Self::Input) -> Self::Output { + self.f.apply(self.g.apply(x).borrow()) + } + + fn domain(&self) -> Option<&dyn Domain> { + self.g.domain() + } + + fn codomain(&self) -> Option<&dyn Domain> { + self.f.codomain() + } +} + +impl Boxed for Compose +where + G::Output: Borrow, +{ + fn hrect(&self) -> &HRect { + self.g.hrect() + } +} + +/// The extension of `f` to multisets, written `f̄` in the paper: it maps a +/// multiset to the multiset of the images of its elements, repeats included. +pub struct MultisetLift { + f: F, +} + +pub fn multiset_lift(f: F) -> MultisetLift +where + F::Input: ToOwned, +{ + MultisetLift { f } +} + +// The elements are taken owned, so that a tensor, which takes `[usize]`, lifts to +// take the `Multiset>` an index function returns. +impl Function for MultisetLift +where + F::Input: ToOwned, +{ + type Input = Multiset<::Owned>; + type Output = Multiset; + + fn apply(&self, x: &Self::Input) -> Self::Output { + x.iter() + .map(|element| self.f.apply(element.borrow())) + .collect() + } + + fn domain(&self) -> Option<&dyn Domain> { + self.f.domain() + } + + fn codomain(&self) -> Option<&dyn Domain> { + self.f.codomain() + } +} + +/// A function on values, given as a closure, from `A` to `B`. Composed after a +/// tensor, it applies element-wise. +pub struct ValueFunction { + // Held as a closure rather than a plain `fn` pointer so that a function can + // carry values of its own, as `scale` and `clamp` do. + f: Rc B>, +} + +impl ValueFunction { + pub fn new(f: impl Fn(A) -> B + 'static) -> Self { + Self { f: Rc::new(f) } + } +} + +impl Function for ValueFunction { + type Input = A; + type Output = B; + + fn apply(&self, x: &A) -> B { + (self.f)(x.clone()) + } +} diff --git a/src/functions.rs b/src/functions.rs index 5e4569a..114b8fc 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -1,52 +1,52 @@ //! Named constructors for the compositions the library can already express. -//! Everything here is assembled from `mapping` and `tensor`; none of it adds -//! machinery of its own. +//! Everything here is assembled from `function`, `mapping` and `tensor`; none +//! of it adds 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; -use crate::domain::HRect; +use crate::domain::{Domain, HRect}; +use crate::function::{Compose, ValueFunction, compose, multiset_lift}; use crate::mapping::{ Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Permutation, Projection, Squeeze, Unsqueeze, Where, }; use crate::tensor::{ - JaggedTensor, LeftComposition, Physical, Product, ProductReduction, Reduction, ReductionTensor, - RightComposition, Sum, SumReduction, Tensor, + BoxedTensor, JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, + ReductionTensor, Sum, SumReduction, Tensor, }; -pub fn square + Copy + 'static, U: Tensor + ?Sized>( - u: Rc, -) -> LeftComposition { - LeftComposition::new(u, |x| x * x) +pub fn square + Copy + 'static, U: Tensor>( + u: U, +) -> Compose, U> { + compose(ValueFunction::new(|x: A| x * x), u) } -pub fn negate + 'static, U: Tensor + ?Sized>( - u: Rc, -) -> LeftComposition { - LeftComposition::new(u, |x| -x) +pub fn negate + Clone + 'static, U: Tensor>( + u: U, +) -> Compose, U> { + compose(ValueFunction::new(|x: A| -x), u) } -pub fn scale + Copy + 'static, U: Tensor + ?Sized>( - u: Rc, +pub fn scale + Copy + 'static, U: Tensor>( + u: U, factor: A, -) -> LeftComposition { - LeftComposition::new(u, move |x| x * factor) +) -> Compose, U> { + compose(ValueFunction::new(move |x: A| x * factor), u) } /// 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 + ?Sized>( - u: Rc, +pub fn clamp>( + u: U, low: A, high: A, -) -> LeftComposition { +) -> Compose, U> { assert!(low <= high, "clamp needs low <= high, got {low} and {high}"); - LeftComposition::new(u, move |x| { + let clamped = ValueFunction::new(move |x: A| { if x < low { low } else if x > high { @@ -54,26 +54,22 @@ pub fn clamp + ?Sized>( } else { x } - }) + }); + compose(clamped, u) } /// Reorders the axes of a tensor: `axes[i]` is the axis of `u` that becomes /// axis `i` of the result. -pub fn permute + 'static>( - u: Rc, - axes: Vec, -) -> RightComposition { +pub fn permute>(u: U, axes: Vec) -> Compose { let codomain = u.hrect().clone(); - RightComposition::new(u, Rc::new(Permutation::new(codomain, axes))) + compose(u, 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 + 'static>( - u: Rc, -) -> RightComposition { - let ndim = u.domain().ndim(); +pub fn transpose>(u: U) -> Compose { + let ndim = u.hrect().ndim(); assert!(ndim >= 2, "transpose needs at least two dimensions"); let mut axes: Vec = (0..ndim).collect(); axes.swap(ndim - 2, ndim - 1); @@ -82,12 +78,12 @@ pub fn transpose + 'static>( /// Moves the axis at `source` to `destination`; the axes left over keep their /// relative order. -pub fn movedim + 'static>( - u: Rc, +pub fn movedim>( + u: U, source: usize, destination: usize, -) -> RightComposition { - let ndim = u.domain().ndim(); +) -> Compose { + let ndim = u.hrect().ndim(); assert!(source < ndim, "source axis {source} is out of range"); assert!( destination < ndim, @@ -101,12 +97,9 @@ pub fn movedim + '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 + 'static>( - u: Rc, - dim: usize, -) -> RightComposition { +pub fn unsqueeze>(u: U, dim: usize) -> Compose { let codomain = u.hrect().clone(); - RightComposition::new(u, Rc::new(Unsqueeze::new(codomain, dim))) + compose(u, Unsqueeze::new(codomain, dim)) } /// The axes a function is told to act on, in place of torch's optional `dim` @@ -149,10 +142,7 @@ 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 + 'static>( - u: Rc, - dims: impl Dims, -) -> RightComposition { +pub fn squeeze>(u: U, dims: impl Dims) -> Compose { let codomain = u.hrect().clone(); let ndim = codomain.shape().len(); let named = dims.axes(ndim); @@ -167,18 +157,18 @@ pub fn squeeze + 'static>( .into_iter() .filter(|&axis| codomain.shape()[axis] == 1) .collect(); - RightComposition::new(u, Rc::new(Squeeze::new(codomain, removed))) + compose(u, Squeeze::new(codomain, removed)) } /// Joins tensors end to end along `dim`; their other extents must agree. -pub fn concatenate( +pub fn concatenate( tensors: Vec>>, dim: usize, -) -> RightComposition { +) -> Compose, Concat> { assert!(!tensors.is_empty(), "concatenate needs at least one tensor"); let underlying = JaggedTensor::new(tensors); let codomain = underlying.jagged_domain().clone(); - RightComposition::new(Rc::new(underlying), Rc::new(Concat::new(codomain, dim))) + compose(underlying, Concat::new(codomain, dim)) } /// Joins tensors of one shape along a new axis at `dim`, by concatenating them @@ -186,7 +176,7 @@ pub fn concatenate( pub fn stack( tensors: Vec>>, dim: usize, -) -> RightComposition { +) -> Compose, Concat> { assert!(!tensors.is_empty(), "stack needs at least one tensor"); let shape = tensors[0].hrect().clone(); assert!( @@ -197,7 +187,7 @@ pub fn stack( .into_iter() .map(|tensor| { let unsqueeze = Unsqueeze::new(tensor.hrect().clone(), dim); - Rc::new(RightComposition::new(tensor, Rc::new(unsqueeze))) as Rc> + Rc::new(compose(tensor, unsqueeze)) as Rc> }) .collect(); concatenate(unsqueezed, dim) @@ -205,61 +195,70 @@ pub fn stack( /// 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_( +pub fn where_( condition: Rc>, input: Rc>, other: Rc>, -) -> RightComposition { +) -> Compose, Where> { let underlying = JaggedTensor::new(vec![input, other]); let codomain = underlying.jagged_domain().clone(); - RightComposition::new( - Rc::new(underlying), - Rc::new(Where::new(codomain, condition)), - ) + compose(underlying, Where::new(codomain, condition)) } /// 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 + 'static>( - u: Rc, +pub fn diagonalize>( + u: U, rank: usize, -) -> Sum +) -> ReductionTensor, U, Diagonal> where SumReduction: Reduction, { - assert_eq!(u.domain().ndim(), 1, "diagonalize takes a vector"); + assert_eq!(u.hrect().ndim(), 1, "diagonalize takes a vector"); assert!(rank > 0, "a diagonal tensor has at least one dimension"); let extent = u.hrect().shape()[0]; - Sum::new(u, Rc::new(Diagonal::new(HRect::new(vec![extent; rank])))) + let diagonal = Diagonal::new(HRect::new(vec![extent; rank])); + compose(Sum::new(), compose(multiset_lift(u), diagonal)) } /// 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>(shape: Vec) -> ReductionTensor { +fn neutral>( + shape: Vec, +) -> ReductionTensor, Physical, EmptyMapping> { let mapping = EmptyMapping::new(HRect::new(shape), HRect::new(vec![0])); - ReductionTensor::new(Rc::new(Physical::new(vec![])), Rc::new(mapping)) + compose( + Reduce::new(), + compose(multiset_lift(Physical::new(vec![])), mapping), + ) } -pub fn zeros(shape: Vec) -> Sum { +pub fn zeros(shape: Vec) -> ReductionTensor { neutral(shape) } -pub fn ones(shape: Vec) -> Product { +pub fn ones(shape: Vec) -> ReductionTensor { neutral(shape) } /// Ones down the diagonal and zeros elsewhere, over no stored data: the ones /// are the neutral element of a product, the zeros that of a sum. -pub fn eye_rect(rows: usize, cols: usize) -> Sum { +pub fn eye_rect( + rows: usize, + cols: usize, +) -> ReductionTensor, Diagonal> { let extent = rows.min(cols); - Sum::new( - Rc::new(ones(vec![extent])), - Rc::new(Diagonal::new(HRect::new(vec![rows, cols]))), + let diagonal = Diagonal::new(HRect::new(vec![rows, cols])); + compose( + Sum::new(), + compose(multiset_lift(ones(vec![extent])), diagonal), ) } -pub fn eye(n: usize) -> Sum { +pub fn eye( + n: usize, +) -> ReductionTensor, Diagonal> { eye_rect(n, n) } @@ -267,11 +266,11 @@ pub fn eye(n: usize) -> Sum { /// `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>( indices: Vec>, - values: Rc>, + values: U, size: Vec, -) -> Sum +) -> ReductionTensor, U, COO> where SumReduction: Reduction, { @@ -281,7 +280,7 @@ where "sparse_coo_tensor needs one row of indices per dimension of size" ); assert_eq!( - values.domain().ndim(), + values.hrect().ndim(), 1, "sparse_coo_tensor takes a vector of values" ); @@ -305,52 +304,60 @@ where "coordinate {coordinate:?} lies outside size {size:?}" ); } - Sum::new(values, Rc::new(COO::new(HRect::new(size), coordinates))) + let coo = COO::new(HRect::new(size), coordinates); + compose(Sum::new(), compose(multiset_lift(values), coo)) } /// A sparse latticed tensor of shape `shape`: the value at index `x` of /// `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 + 'static>( - physical: Rc, +pub fn sparse_latticed_tensor>( + physical: U, basis: Vec>, offset: Vec, shape: Vec, -) -> Sum +) -> ReductionTensor, U, InvAffine> where SumReduction: Reduction, { let affine = Affine::new(physical.hrect().clone(), basis, offset, HRect::new(shape)); - Sum::new(physical, Rc::new(InvAffine::new(affine))) + compose( + Sum::new(), + compose(multiset_lift(physical), InvAffine::new(affine)), + ) } -pub fn sum + 'static>( - input: Rc, +pub fn sum>( + input: U, dim: usize, -) -> Sum +) -> ReductionTensor, U, InvProjection> where SumReduction: Reduction, { - let codomain = input.hrect().clone(); - Sum::new(input, Rc::new(InvProjection::new(codomain, dim))) + let rows = InvProjection::new(input.hrect().clone(), dim); + compose(Sum::new(), compose(multiset_lift(input), rows)) } -pub fn prod + 'static>( - input: Rc, +pub fn prod>( + input: U, dim: usize, -) -> Product +) -> ReductionTensor, U, InvProjection> where ProductReduction: Reduction, { - let codomain = input.hrect().clone(); - Product::new(input, Rc::new(InvProjection::new(codomain, dim))) + let rows = InvProjection::new(input.hrect().clone(), dim); + compose(Product::new(), compose(multiset_lift(input), rows)) } -pub fn mm + 'static, V: BoxedTensor + 'static>( - mat1: Rc, - mat2: Rc, -) -> Sum +/// The tensors `mm` multiplies and then sums, stacked: the product of the two +/// matrices at each `[m, n, p]`. +type Products = ReductionTensor, Compose, Concat>, InvProjection>; + +pub fn mm + 'static, V: BoxedTensor + 'static>( + mat1: U, + mat2: V, +) -> ReductionTensor, Products, InvProjection> where SumReduction: Reduction, ProductReduction: Reduction, @@ -369,16 +376,21 @@ where let latent_domain = HRect::new(vec![m, n, p]); let stacked_latent_domain = HRect::new(vec![2, m, n, p]); - let phi1 = Rc::new(Projection::new(latent_domain.clone(), 2)); // phi1(m, n, p) = [m, n]^T - let phi2 = Rc::new(Projection::new(latent_domain.clone(), 0)); // phi2(m, n, p) = [n, p]^T - let latent_inv_proj = Rc::new(InvProjection::new(stacked_latent_domain, 0)); // l(2, m, n, p) = [m, n, p]^T - let psi_inv = Rc::new(InvProjection::new(latent_domain.clone(), 1)); // psi(m, n, p) = [m, p]^T + let phi1 = Projection::new(latent_domain.clone(), 2); // phi1(m, n, p) = [m, n]^T + let phi2 = Projection::new(latent_domain.clone(), 0); // phi2(m, n, p) = [n, p]^T + let latent_inv_proj = InvProjection::new(stacked_latent_domain, 0); // l(2, m, n, p) = [m, n, p]^T + let psi_inv = InvProjection::new(latent_domain.clone(), 1); // psi(m, n, p) = [m, p]^T - let latent_mat1 = Rc::new(RightComposition::new(mat1, phi1)); // shape [m, n, p] - let latent_mat2 = Rc::new(RightComposition::new(mat2, phi2)); // shape [m, n, p] - let stacked = Rc::new(stack(vec![latent_mat1, latent_mat2], 0)); // shape [2, m, n, p] + let latent_mats: Vec>> = vec![ + Rc::new(compose(mat1, phi1)), // shape [m, n, p] + Rc::new(compose(mat2, phi2)), // shape [m, n, p] + ]; + let stacked = stack(latent_mats, 0); // shape [2, m, n, p] - let prod_tensor = Rc::new(Product::new(stacked, latent_inv_proj)); // shape [m, n, p] + let prod_tensor = compose( + Product::new(), + compose(multiset_lift(stacked), latent_inv_proj), + ); // shape [m, n, p] - Sum::new(prod_tensor, psi_inv) // shape [m, p] + compose(Sum::new(), compose(multiset_lift(prod_tensor), psi_inv)) // shape [m, p] } diff --git a/src/lib.rs b/src/lib.rs index a88e64f..6ff1837 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod display; pub mod domain; +pub mod function; pub mod functions; mod lattice; pub mod mapping; @@ -10,22 +11,21 @@ pub mod multiset; pub mod signature; pub mod tensor; -pub use display::{BoxedTensor, format_boxed}; +pub use display::format_boxed; pub use domain::{Domain, HRect, JaggedDomain}; +pub use function::{Boxed, Compose, Function, MultisetLift, ValueFunction, compose, multiset_lift}; pub use functions::{ 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::{ - Affine, BoxedMapping, BoxedPowersetMapping, COO, Concat, Diagonal, EmptyMapping, IndexMapping, - InvAffine, InvProjection, Permutation, PowersetMapping, Projection, Squeeze, Stride, Unsqueeze, - Where, + Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Mapping, Permutation, + Projection, Squeeze, Stride, Unsqueeze, Where, }; pub use multiset::Multiset; pub use signature::{HRectSignature, Signature}; pub use tensor::{ - Bar, BarComposition, Inf, InfReduction, JaggedTensor, LeftComposition, Physical, Product, - ProductReduction, Reduction, ReductionTensor, RightComposition, Sum, SumReduction, Sup, - SupReduction, Tensor, + BoxedTensor, Inf, InfReduction, JaggedTensor, Physical, Product, ProductReduction, Reduce, + Reduction, ReductionTensor, Sum, SumReduction, Sup, SupReduction, Tensor, }; diff --git a/src/main.rs b/src/main.rs index b8d88bc..1b89d3b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - COO, HRect, Physical, Product, RightComposition, Stride, Sum, Sup, Tensor, diagonalize, eye, - movedim, negate, ones, square, transpose, zeros, + COO, Function, HRect, Physical, Product, Stride, Sum, Sup, compose, diagonalize, eye, movedim, + multiset_lift, negate, ones, square, transpose, zeros, }; fn main() { @@ -27,12 +27,16 @@ fn main() { // Tensors let physical = Rc::new(Physical::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])); - let strided = Rc::new(RightComposition::new(physical.clone(), stride)); - let sparse_coo_vector = Rc::new(Sum::new(physical.clone(), coo_vector)); - let repeated_sparse = Rc::new(RightComposition::new(sparse_coo_vector.clone(), repeat)); - let sparse_coo = Rc::new(Sum::new(physical.clone(), coo.clone())); - let prod_coo = Rc::new(Product::new(physical.clone(), coo.clone())); - let sup_coo = Rc::new(Sup::new(physical.clone(), coo)); + let lifted = Rc::new(multiset_lift(physical.clone())); + let strided = Rc::new(compose(physical.clone(), stride)); + let sparse_coo_vector = Rc::new(compose(Sum::new(), compose(lifted.clone(), coo_vector))); + let repeated_sparse = Rc::new(compose(sparse_coo_vector.clone(), repeat)); + let sparse_coo = Rc::new(compose(Sum::new(), compose(lifted.clone(), coo.clone()))); + let prod_coo = Rc::new(compose( + Product::new(), + compose(lifted.clone(), coo.clone()), + )); + let sup_coo = Rc::new(compose(Sup::new(), compose(lifted, coo))); let diag2d = Rc::new(diagonalize(physical.clone(), 2)); let diag3d = Rc::new(diagonalize(physical.clone(), 3)); let identity = Rc::new(eye(6)); @@ -64,7 +68,7 @@ fn main() { "strided({}, {}) = {}", index[0], index[1], - strided.get(&index) + strided.apply(&index) ); println!("\ntransposed:\n{transposed}"); diff --git a/src/mapping.rs b/src/mapping.rs index 75e36f1..9507fa7 100644 --- a/src/mapping.rs +++ b/src/mapping.rs @@ -1,99 +1,44 @@ use std::rc::Rc; -use crate::display::BoxedTensor; use crate::domain::{Domain, HRect, JaggedDomain}; +use crate::function::{Boxed, Function}; use crate::lattice::Preimage; use crate::multiset::Multiset; use crate::signature::{HRectSignature, Signature}; - -/// A mapping between two index sets, from its domain to its codomain. -pub trait IndexMapping { - fn map(&self, index: &[usize]) -> Vec; - - fn domain(&self) -> &dyn Domain; - - fn codomain(&self) -> &dyn Domain; -} - -/// A mapping whose domain is a hrect. Composing a tensor over one of these gives -/// a boxed tensor, hence a printable one. -pub trait BoxedMapping: IndexMapping { - fn hrect(&self) -> &HRect; -} - -/// A mapping sending each index of its domain to a multiset of indices of its -/// codomain, an index function in the paper's terms. Note that the codomain is -/// the set the returned indices are drawn from, not its multisets. -pub trait PowersetMapping { - fn map(&self, index: &[usize]) -> Multiset>; - - fn domain(&self) -> &dyn Domain; - - fn codomain(&self) -> &dyn Domain; -} - -/// A powerset mapping whose domain is a hrect. Reducing over one of these gives a -/// boxed tensor, hence a printable one. -pub trait BoxedPowersetMapping: PowersetMapping { - fn hrect(&self) -> &HRect; -} - -/// What a mapping from a hyperrectangle has to provide. Implementing this rather -/// than `IndexMapping` gets the domain accessors and printability below for -/// free, whatever the codomain; a mapping over a domain that is not a hrect -/// implements `IndexMapping` directly instead, and is not printable. -pub trait HRectMapping { +use crate::tensor::BoxedTensor; + +/// A function on the indices of a hrect: an index map when `Output` is an index +/// of its codomain, and an index function when it is a multiset of them. The +/// codomain is the set those indices are drawn from. Implementing this gets +/// `Function` and `Boxed` for free, with the domain and codomain of the +/// signature. +pub trait Mapping { + type Output; type Codomain: Domain; - fn map(&self, index: &[usize]) -> Vec; + fn map(&self, index: &[usize]) -> Self::Output; fn signature(&self) -> &Signature; } -impl IndexMapping for M { - fn map(&self, index: &[usize]) -> Vec { - HRectMapping::map(self, index) - } - - fn domain(&self) -> &dyn Domain { - self.signature().domain() - } - - fn codomain(&self) -> &dyn Domain { - self.signature().codomain() - } -} - -impl BoxedMapping for M { - fn hrect(&self) -> &HRect { - self.signature().domain() - } -} +impl Function for M { + type Input = [usize]; + type Output = M::Output; -/// The powerset counterpart of `HRectMapping`. -pub trait HRectPowersetMapping { - type Codomain: Domain; - - fn map(&self, index: &[usize]) -> Multiset>; - - fn signature(&self) -> &Signature; -} - -impl PowersetMapping for M { - fn map(&self, index: &[usize]) -> Multiset> { - HRectPowersetMapping::map(self, index) + fn apply(&self, index: &[usize]) -> M::Output { + self.map(index) } - fn domain(&self) -> &dyn Domain { - self.signature().domain() + fn domain(&self) -> Option<&dyn Domain> { + Some(self.signature().domain()) } - fn codomain(&self) -> &dyn Domain { - self.signature().codomain() + fn codomain(&self) -> Option<&dyn Domain> { + Some(self.signature().codomain()) } } -impl BoxedPowersetMapping for M { +impl Boxed for M { fn hrect(&self) -> &HRect { self.signature().domain() } @@ -135,7 +80,8 @@ fn flat_extent(shape: &[usize], strides: &[usize]) -> usize { .sum::() } -impl HRectMapping for Stride { +impl Mapping for Stride { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -185,7 +131,8 @@ impl Permutation { } } -impl HRectMapping for Permutation { +impl Mapping for Permutation { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -224,7 +171,8 @@ impl Unsqueeze { } } -impl HRectMapping for Unsqueeze { +impl Mapping for Unsqueeze { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -271,7 +219,8 @@ impl Squeeze { } } -impl HRectMapping for Squeeze { +impl Mapping for Squeeze { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -325,7 +274,8 @@ impl Concat { } } -impl HRectMapping for Concat { +impl Mapping for Concat { + type Output = Vec; type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { @@ -374,11 +324,12 @@ impl Where { } } -impl HRectMapping for Where { +impl Mapping for Where { + type Output = Vec; type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { - let part = if self.condition.get(index) { 0 } else { 1 }; + let part = if self.condition.apply(index) { 0 } else { 1 }; let mut result = Vec::with_capacity(index.len() + 1); result.push(part); result.extend_from_slice(index); @@ -402,7 +353,8 @@ impl EmptyMapping { } } -impl HRectPowersetMapping for EmptyMapping { +impl Mapping for EmptyMapping { + type Output = Multiset>; type Codomain = HRect; fn map(&self, _index: &[usize]) -> Multiset> { @@ -433,7 +385,8 @@ impl COO { } } -impl HRectPowersetMapping for COO { +impl Mapping for COO { + type Output = Multiset>; type Codomain = HRect; fn map(&self, index: &[usize]) -> Multiset> { @@ -470,7 +423,8 @@ impl Diagonal { } } -impl HRectPowersetMapping for Diagonal { +impl Mapping for Diagonal { + type Output = Multiset>; type Codomain = HRect; fn map(&self, index: &[usize]) -> Multiset> { @@ -513,7 +467,8 @@ impl Projection { } } -impl HRectMapping for Projection { +impl Mapping for Projection { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -551,7 +506,8 @@ impl InvProjection { } } -impl HRectPowersetMapping for InvProjection { +impl Mapping for InvProjection { + type Output = Multiset>; type Codomain = HRect; fn map(&self, index: &[usize]) -> Multiset> { @@ -623,7 +579,8 @@ impl Affine { } } -impl HRectMapping for Affine { +impl Mapping for Affine { + type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -679,7 +636,8 @@ impl InvAffine { } } -impl HRectPowersetMapping for InvAffine { +impl Mapping for InvAffine { + type Output = Multiset>; type Codomain = HRect; fn map(&self, index: &[usize]) -> Multiset> { diff --git a/src/tensor.rs b/src/tensor.rs index 6d83f42..8e16513 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,33 +1,21 @@ use std::marker::PhantomData; use std::rc::Rc; -use crate::display::BoxedTensor; use crate::domain::{Domain, HRect, JaggedDomain}; -use crate::mapping::{BoxedMapping, BoxedPowersetMapping, IndexMapping, PowersetMapping}; +use crate::function::{Boxed, Compose, Function, MultisetLift}; use crate::multiset::Multiset; -/// A mapping from an index set to a value set. The index set is the tensor's -/// domain, and the value set is its value type `A`, `f32` unless said -/// otherwise. -pub trait Tensor { - fn get(&self, index: &[usize]) -> A; +/// A function from indices to values, with values of type `A`, `f32` unless +/// said otherwise. Nothing implements this directly: every function of that +/// type is a tensor. +pub trait Tensor: Function {} - fn domain(&self) -> &dyn Domain; -} +impl + ?Sized> Tensor for T {} -/// Checks that a mapping lands on the indices the underlying tensor reads, -/// which is all that can be checked without enumerating a domain. In -/// particular, a mapping into flat positions needs a one-dimensional tensor -/// under it. -fn check_composable(codomain: &dyn Domain, u: &dyn Tensor) { - assert_eq!( - codomain.ndim(), - u.domain().ndim(), - "the mapping lands on {}-dimensional indices, but the underlying tensor is indexed by {} coordinates", - codomain.ndim(), - u.domain().ndim() - ); -} +/// A tensor whose domain is a hrect, hence a printable one. +pub trait BoxedTensor: Tensor + Boxed {} + +impl + Boxed + ?Sized> BoxedTensor for T {} pub struct Physical { values: Vec, @@ -41,8 +29,11 @@ impl Physical { } } -impl Tensor for Physical { - fn get(&self, index: &[usize]) -> A { +impl Function for Physical { + type Input = [usize]; + type Output = A; + + fn apply(&self, index: &[usize]) -> A { assert_eq!( index.len(), 1, @@ -51,20 +42,20 @@ impl Tensor for Physical { self.values[index[0]].clone() } - fn domain(&self) -> &dyn Domain { - &self.domain + fn domain(&self) -> Option<&dyn Domain> { + Some(&self.domain) } } -impl BoxedTensor for Physical { +impl Boxed for Physical { fn hrect(&self) -> &HRect { &self.domain } } /// Several tensors read as one: `[t, rest..]` reads part `t` at `rest`. This is -/// the underlying that concatenating, stacking and `where` remap into. Its domain -/// is jagged rather than a hrect, so it is not printable. +/// the paper's stack, which concatenating, stacking and `where` remap into. Its +/// domain is jagged rather than a hrect, so it is not printable. pub struct JaggedTensor { parts: Vec>>, domain: JaggedDomain, @@ -81,94 +72,23 @@ impl JaggedTensor { } } -impl Tensor for JaggedTensor { - fn get(&self, index: &[usize]) -> A { - self.parts[index[0]].get(&index[1..]) - } - - fn domain(&self) -> &dyn Domain { - &self.domain - } -} - -// The domain is the mapping's own domain, so it is never passed in. The mapping -// type is a parameter rather than always `dyn IndexMapping` so that composing -// over a boxed mapping is known to give a boxed, printable tensor. Pass -// `dyn BoxedMapping` or `dyn IndexMapping` for it to erase the mapping. The -// values are the underlying tensor's, read elsewhere. -pub struct RightComposition { - u: Rc>, - mapping: Rc, -} - -impl RightComposition { - pub fn new(u: Rc>, mapping: Rc) -> Self { - check_composable(mapping.codomain(), &*u); - Self { u, mapping } - } -} +impl Function for JaggedTensor { + type Input = [usize]; + type Output = A; -impl Tensor for RightComposition { - fn get(&self, virtual_index: &[usize]) -> A { - let underlying_index = self.mapping.map(virtual_index); - self.u.get(&underlying_index) + fn apply(&self, index: &[usize]) -> A { + self.parts[index[0]].apply(&index[1..]) } - fn domain(&self) -> &dyn Domain { - self.mapping.domain() + fn domain(&self) -> Option<&dyn Domain> { + Some(&self.domain) } } -impl BoxedTensor for RightComposition { - fn hrect(&self) -> &HRect { - self.mapping.hrect() - } -} - -// Transforming the values leaves the index set alone, so the domain is the -// underlying tensor's. Its type is a parameter so that a transform of a boxed -// tensor is known to be boxed, and printable, in turn. The transform may change -// the value type, from `A` to `B`, as a comparison turns numbers into booleans. -pub struct LeftComposition { - u: Rc, - // Held as a closure rather than a plain `fn` pointer so that a transform can - // carry values of its own, as `scale` and `clamp` do. - f: Rc B>, -} - -impl + ?Sized, A, B> LeftComposition { - pub fn new(u: Rc, f: impl Fn(A) -> B + 'static) -> Self { - Self { u, f: Rc::new(f) } - } -} - -impl + ?Sized, A, B> Tensor for LeftComposition { - fn get(&self, index: &[usize]) -> B { - (self.f)(self.u.get(index)) - } - - fn domain(&self) -> &dyn Domain { - self.u.domain() - } -} - -impl + ?Sized, A, B> BoxedTensor for LeftComposition { - fn hrect(&self) -> &HRect { - self.u.hrect() - } -} - -/// A way of combining the values of a tensor of value type `A` into one. +/// A way of combining values of type `A` into one, from a neutral element. pub trait Reduction { const NEUTRAL: A; fn combine(a: A, b: A) -> A; - - /// Reduces a multiset to a single value: the neutral element for the empty - /// multiset, its element for a multiset of one, and otherwise its elements - /// combined in any order, as `combine` is associative and commutative. - fn reduce(multiset: Multiset) -> A { - multiset.into_iter().fold(Self::NEUTRAL, Self::combine) - } } pub struct SumReduction; @@ -216,94 +136,45 @@ reductions! { i64: 0, 1, i64::MIN, i64::MAX; } -/// The extension of a tensor `U` to multisets, written `Ū`: it maps a multiset of -/// indices to the multiset of the values of `U` at them, repeats included. Its -/// indices being multisets, it is not a tensor itself; composed with an index -/// function, it gives one. -pub struct Bar { - u: Rc>, +/// The reduction `R` as a function from multisets of values to values, `Ω` in +/// the paper: the neutral element for the empty multiset, the element of a +/// multiset of one, and otherwise its elements combined in any order, as +/// `combine` is associative and commutative. +pub struct Reduce { + _marker: PhantomData (R, A)>, } -impl Bar { - pub fn new(u: Rc>) -> Self { - Self { u } - } - - pub fn get(&self, indices: &Multiset>) -> Multiset { - indices.iter().map(|index| self.u.get(index)).collect() - } -} - -// `Ū ∘ ψ`: at each index, the multiset of the values of `U` at the indices `ψ` -// maps it to. This is the multiset-valued tensor a reduction reduces. As in a -// right composition, the domain is the mapping's own domain, and the mapping -// type is a parameter so that a boxed mapping is known to give a boxed, -// printable tensor. -pub struct BarComposition { - bar: Bar, - mapping: Rc, -} - -impl BarComposition { - pub fn new(u: Rc>, mapping: Rc) -> Self { - check_composable(mapping.codomain(), &*u); +impl Reduce { + pub fn new() -> Self { Self { - bar: Bar::new(u), - mapping, + _marker: PhantomData, } } } -impl Tensor> for BarComposition { - fn get(&self, index: &[usize]) -> Multiset { - self.bar.get(&self.mapping.map(index)) - } - - fn domain(&self) -> &dyn Domain { - self.mapping.domain() - } -} - -impl BoxedTensor> for BarComposition { - fn hrect(&self) -> &HRect { - self.mapping.hrect() +// Written out rather than derived, which would require `R: Default` and `A: Default`. +impl Default for Reduce { + fn default() -> Self { + Self::new() } } -// `Ω ∘ Ū ∘ ψ`: the reduction `R` of the multisets that a bar composition gathers. -pub struct ReductionTensor { - gathered: BarComposition, - _marker: PhantomData R>, -} +impl, A: Clone> Function for Reduce { + type Input = Multiset; + type Output = A; -impl, M: PowersetMapping + ?Sized, A> ReductionTensor { - pub fn new(u: Rc>, mapping: Rc) -> Self { - Self { - gathered: BarComposition::new(u, mapping), - _marker: PhantomData, - } + fn apply(&self, multiset: &Multiset) -> A { + multiset.iter().cloned().fold(R::NEUTRAL, R::combine) } } -impl, M: PowersetMapping + ?Sized, A> Tensor for ReductionTensor { - fn get(&self, index: &[usize]) -> A { - R::reduce(self.gathered.get(index)) - } - - fn domain(&self) -> &dyn Domain { - self.gathered.domain() - } -} - -impl, M: BoxedPowersetMapping + ?Sized, A> BoxedTensor - for ReductionTensor -{ - fn hrect(&self) -> &HRect { - self.gathered.hrect() - } -} +pub type Sum = Reduce; +pub type Product = Reduce; +pub type Sup = Reduce; +pub type Inf = Reduce; -pub type Sum = ReductionTensor; -pub type Product = ReductionTensor; -pub type Sup = ReductionTensor; -pub type Inf = ReductionTensor; +/// `Ω ∘ Ū ∘ ψ`, the reduction tensor of the paper, with reduction `R`, +/// underlying tensor `U` and index function `M`: at each index, the reduction +/// of the multiset of the values of `U` at the indices `M` maps it to. Built +/// as `compose(Sum::new(), compose(multiset_lift(u), psi))`. +pub type ReductionTensor = Compose, M>>; diff --git a/tests/display.rs b/tests/display.rs index b28fe8d..78ef74d 100644 --- a/tests/display.rs +++ b/tests/display.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - BarComposition, COO, Diagonal, HRect, Physical, RightComposition, Stride, Sum, Sup, Tensor, - format_boxed, + COO, Diagonal, Function, HRect, Physical, Stride, Sum, Sup, compose, format_boxed, + multiset_lift, }; fn physical() -> Rc { @@ -27,10 +27,10 @@ fn tensors_of_other_value_types_print_the_same_way() { // its multisets of values are summed. #[test] fn a_multiset_valued_tensor_prints_its_multisets() { - let stored = Rc::new(Physical::new(vec![1.0, 2.0, 4.0])); - let diagonal = Rc::new(Diagonal::new(HRect::new(vec![3, 3]))); - let before = BarComposition::new(stored.clone(), diagonal.clone()); - let after = Sum::new(stored, diagonal); + let stored = Physical::new(vec![1.0, 2.0, 4.0]); + let diagonal = Diagonal::new(HRect::new(vec![3, 3])); + let before = Rc::new(compose(multiset_lift(stored), diagonal)); + let after = compose(Sum::new(), before.clone()); assert_eq!( before.to_string(), @@ -44,7 +44,7 @@ fn a_multiset_valued_tensor_prints_its_multisets() { #[test] fn a_rank_two_tensor_prints_one_row_per_line_aligned_under_the_bracket() { - let strided = RightComposition::new( + let strided = compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), ); @@ -58,7 +58,7 @@ fn a_rank_two_tensor_prints_one_row_per_line_aligned_under_the_bracket() { #[test] fn values_are_right_aligned_to_a_common_width() { let wide = Rc::new(Physical::new(vec![1.0, 16.0, 4.0, 25.0])); - let strided = RightComposition::new( + let strided = compose( wide, Rc::new(Stride::new(HRect::new(vec![2, 2]), vec![1, 2])), ); @@ -72,7 +72,7 @@ fn the_neutral_element_of_a_sup_prints_as_negative_infinity() { HRect::new(vec![2, 2]), vec![vec![0, 0], vec![1, 1]], )); - let sup = Sup::new(physical(), coo); + let sup = compose(Sup::new(), compose(multiset_lift(physical()), coo)); assert_eq!( sup.to_string(), @@ -82,9 +82,12 @@ fn the_neutral_element_of_a_sup_prints_as_negative_infinity() { #[test] fn rank_three_tensors_separate_blocks_with_a_blank_line() { - let diag = Sum::new( - physical(), - Rc::new(Diagonal::new(HRect::new(vec![2, 2, 2]))), + let diag = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(Diagonal::new(HRect::new(vec![2, 2, 2]))), + ), ); assert_eq!( @@ -101,18 +104,21 @@ fn rank_three_tensors_separate_blocks_with_a_blank_line() { #[test] fn display_agrees_with_direct_lookups() { - let sparse = Sum::new( - physical(), - Rc::new(COO::new( - HRect::new(vec![2, 3]), - vec![vec![0, 0], vec![1, 2]], - )), + let sparse = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(COO::new( + HRect::new(vec![2, 3]), + vec![vec![0, 0], vec![1, 2]], + )), + ), ); let rendered = sparse.to_string(); for i in 0..2 { for j in 0..3 { - let value = sparse.get(&[i, j]); + let value = sparse.apply(&[i, j]); assert!( rendered.contains(&value.to_string()), "{rendered} is missing the value at ({i}, {j})" @@ -123,7 +129,7 @@ fn display_agrees_with_direct_lookups() { #[test] fn boxed_tensors_print_through_their_hrect() { - let strided = RightComposition::new( + let strided = compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), ); diff --git a/tests/functions.rs b/tests/functions.rs index 873b07e..c52f9bd 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -1,18 +1,20 @@ use std::rc::Rc; use cten::{ - BoxedTensor, COO, Domain, HRect, LeftComposition, Physical, RightComposition, Stride, Sum, - Tensor, 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, + Boxed, BoxedTensor, COO, Compose, Domain, Function, HRect, Physical, Stride, Sum, Tensor, + ValueFunction, clamp, compose, concatenate, diagonalize, eye, eye_rect, mm, movedim, + multiset_lift, negate, ones, permute, prod, scale, sparse_coo_tensor, sparse_latticed_tensor, + square, squeeze, stack, sum, transpose, unsqueeze, where_, zeros, }; +type Strided = Compose>, Rc>; + fn physical() -> Rc { Rc::new(Physical::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) } -fn strided() -> Rc> { - Rc::new(RightComposition::new( +fn strided() -> Rc { + Rc::new(compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), )) @@ -20,37 +22,37 @@ fn strided() -> Rc> { /// A rank-three tensor whose values are all distinct, so a permutation of its /// axes is checkable: `[i, j, k]` holds `12i + 4j + k`. -fn cube() -> Rc> { +fn cube() -> Rc { let values: Vec = (0..24).map(|value| value as f32).collect(); - Rc::new(RightComposition::new( + Rc::new(compose( Rc::new(Physical::new(values)), Rc::new(Stride::new(HRect::new(vec![2, 3, 4]), vec![12, 4, 1])), )) } /// A row-major matrix over stored `values`. -fn matrix(rows: usize, cols: usize, values: Vec) -> Rc> { - Rc::new(RightComposition::new( +fn matrix(rows: usize, cols: usize, values: Vec) -> Rc { + Rc::new(compose( Rc::new(Physical::new(values)), Rc::new(Stride::new(HRect::new(vec![rows, cols]), vec![cols, 1])), )) } /// A row-major matrix over stored integers. -fn int_matrix(rows: usize, cols: usize, values: Vec) -> Rc> { - Rc::new(RightComposition::new( +fn int_matrix(rows: usize, cols: usize, values: Vec) -> Rc> { + Rc::new(compose( Rc::new(Physical::new(values)), Rc::new(Stride::new(HRect::new(vec![rows, cols]), vec![cols, 1])), )) } fn collect_1d(tensor: &dyn Tensor, len: usize) -> Vec { - (0..len).map(|i| tensor.get(&[i])).collect() + (0..len).map(|i| tensor.apply(&[i])).collect() } fn collect_2d(tensor: &dyn Tensor, rows: usize, cols: usize) -> Vec> { (0..rows) - .map(|i| (0..cols).map(|j| tensor.get(&[i, j])).collect()) + .map(|i| (0..cols).map(|j| tensor.apply(&[i, j])).collect()) .collect() } @@ -61,8 +63,8 @@ fn square_and_negate_transform_values() { for i in 0..6 { let value = (i + 1) as f32; - assert_eq!(squared.get(&[i]), value * value); - assert_eq!(negated.get(&[i]), -value); + assert_eq!(squared.apply(&[i]), value * value); + assert_eq!(negated.apply(&[i]), -value); } } @@ -72,7 +74,7 @@ fn elementwise_functions_compose() { for i in 0..6 { let value = (i + 1) as f32; - assert_eq!(squared_negative.get(&[i]), value * value); + assert_eq!(squared_negative.apply(&[i]), value * value); } } @@ -84,13 +86,13 @@ fn parameterised_transforms_remember_their_argument() { for i in 0..6 { let value = (i + 1) as f32; - assert_eq!(doubled.get(&[i]), value * 2.0); + assert_eq!(doubled.apply(&[i]), value * 2.0); } // [1, 2, 3, 4, 5, 6] held between 2 and 4. let expected = [2.0, 2.0, 3.0, 4.0, 4.0, 4.0]; for (i, expected) in expected.iter().enumerate() { - assert_eq!(bounded.get(&[i]), *expected); + assert_eq!(bounded.apply(&[i]), *expected); } } @@ -99,7 +101,7 @@ fn parameterised_transforms_compose() { let scaled_twice = scale(Rc::new(scale(physical(), 2.0)), 3.0); for i in 0..6 { - assert_eq!(scaled_twice.get(&[i]), (i + 1) as f32 * 6.0); + assert_eq!(scaled_twice.apply(&[i]), (i + 1) as f32 * 6.0); } } @@ -109,8 +111,8 @@ fn several_transforms_of_one_tensor_keep_their_own_arguments() { let doubled = scale(shared.clone(), 2.0); let tripled = scale(shared, 3.0); - assert_eq!(doubled.get(&[0]), 2.0); - assert_eq!(tripled.get(&[0]), 3.0); + assert_eq!(doubled.apply(&[0]), 2.0); + assert_eq!(tripled.apply(&[0]), 3.0); } #[test] @@ -158,7 +160,7 @@ fn permute_reorders_axes() { for a in 0..4 { for b in 0..2 { for c in 0..3 { - assert_eq!(permuted.get(&[a, b, c]), cube.get(&[b, c, a])); + assert_eq!(permuted.apply(&[a, b, c]), cube.apply(&[b, c, a])); } } } @@ -185,7 +187,7 @@ fn transpose_leaves_leading_axes_alone() { for i in 0..2 { for j in 0..4 { for k in 0..3 { - assert_eq!(transposed.get(&[i, j, k]), cube.get(&[i, k, j])); + assert_eq!(transposed.apply(&[i, j, k]), cube.apply(&[i, k, j])); } } } @@ -204,23 +206,26 @@ fn transpose_is_an_involution() { // that have no strides to permute in the first place. #[test] fn transpose_applies_to_a_tensor_with_no_strides() { - let sparse = Rc::new(Sum::new( - Rc::new(Physical::new(vec![1.0, 2.0])), - Rc::new(COO::new( - HRect::new(vec![2, 3]), - vec![vec![0, 1], vec![1, 2]], - )), + let sparse = Rc::new(compose( + Sum::new(), + compose( + multiset_lift(Rc::new(Physical::new(vec![1.0, 2.0]))), + Rc::new(COO::new( + HRect::new(vec![2, 3]), + vec![vec![0, 1], vec![1, 2]], + )), + ), )); let transposed = transpose(sparse.clone()); assert_eq!(transposed.hrect().shape(), &[3, 2]); - assert_eq!(transposed.get(&[1, 0]), 1.0); - assert_eq!(transposed.get(&[2, 1]), 2.0); - assert_eq!(transposed.get(&[0, 0]), 0.0); + assert_eq!(transposed.apply(&[1, 0]), 1.0); + assert_eq!(transposed.apply(&[2, 1]), 2.0); + assert_eq!(transposed.apply(&[0, 0]), 0.0); for a in 0..3 { for b in 0..2 { - assert_eq!(transposed.get(&[a, b]), sparse.get(&[b, a])); + assert_eq!(transposed.apply(&[a, b]), sparse.apply(&[b, a])); } } } @@ -236,7 +241,7 @@ fn movedim_follows_the_pytorch_convention() { for a in 0..3 { for b in 0..4 { for c in 0..2 { - assert_eq!(moved.get(&[a, b, c]), cube.get(&[c, a, b])); + assert_eq!(moved.apply(&[a, b, c]), cube.apply(&[c, a, b])); } } } @@ -251,7 +256,7 @@ fn movedim_is_reversible() { for i in 0..2 { for j in 0..3 { for k in 0..4 { - assert_eq!(restored.get(&[i, j, k]), cube.get(&[i, j, k])); + assert_eq!(restored.apply(&[i, j, k]), cube.apply(&[i, j, k])); } } } @@ -265,7 +270,7 @@ fn diagonalize_builds_a_diagonal_matrix() { for i in 0..6 { for j in 0..6 { let expected = if i == j { (i + 1) as f32 } else { 0.0 }; - assert_eq!(diagonal.get(&[i, j]), expected); + assert_eq!(diagonal.apply(&[i, j]), expected); } } } @@ -275,16 +280,16 @@ fn diagonalize_generalises_to_higher_rank() { let diagonal = diagonalize(physical(), 3); assert_eq!(diagonal.hrect().shape(), &[6, 6, 6]); - assert_eq!(diagonal.get(&[2, 2, 2]), 3.0); - assert_eq!(diagonal.get(&[2, 2, 1]), 0.0); - assert_eq!(diagonal.get(&[0, 1, 2]), 0.0); - assert_eq!(diagonal.get(&[5, 5, 5]), 6.0); + assert_eq!(diagonal.apply(&[2, 2, 2]), 3.0); + assert_eq!(diagonal.apply(&[2, 2, 1]), 0.0); + assert_eq!(diagonal.apply(&[0, 1, 2]), 0.0); + assert_eq!(diagonal.apply(&[5, 5, 5]), 6.0); } #[test] fn zeros_and_ones_are_constant_at_any_rank() { - assert_eq!(zeros(vec![4]).get(&[3]), 0.0); - assert_eq!(ones(vec![4]).get(&[3]), 1.0); + assert_eq!(zeros(vec![4]).apply(&[3]), 0.0); + assert_eq!(ones(vec![4]).apply(&[3]), 1.0); let flat_zeros = zeros(vec![2, 3]); let flat_ones = ones(vec![2, 3]); @@ -292,8 +297,8 @@ fn zeros_and_ones_are_constant_at_any_rank() { assert_eq!(collect_2d(&flat_zeros, 2, 3), vec![vec![0.0; 3]; 2]); assert_eq!(collect_2d(&flat_ones, 2, 3), vec![vec![1.0; 3]; 2]); - assert_eq!(ones(vec![2, 3, 4]).get(&[1, 2, 3]), 1.0); - assert_eq!(zeros(vec![2, 3, 4]).get(&[1, 2, 3]), 0.0); + assert_eq!(ones(vec![2, 3, 4]).apply(&[1, 2, 3]), 1.0); + assert_eq!(zeros(vec![2, 3, 4]).apply(&[1, 2, 3]), 0.0); } // A hrect of no dimensions holds the single empty index, so a constant over it is @@ -303,8 +308,8 @@ fn a_rank_zero_constant_is_a_scalar() { let one = ones(vec![]); assert_eq!(one.hrect().ndim(), 0); - assert_eq!(one.get(&[]), 1.0); - assert_eq!(zeros(vec![]).get(&[]), 0.0); + assert_eq!(one.apply(&[]), 1.0); + assert_eq!(zeros(vec![]).apply(&[]), 0.0); } #[test] @@ -315,7 +320,7 @@ fn eye_is_the_identity() { for i in 0..6 { for j in 0..6 { let expected = if i == j { 1.0 } else { 0.0 }; - assert_eq!(identity.get(&[i, j]), expected); + assert_eq!(identity.apply(&[i, j]), expected); } } } @@ -350,7 +355,7 @@ fn eye_survives_a_transpose() { for i in 0..4 { for j in 0..4 { - assert_eq!(transposed.get(&[i, j]), identity.get(&[i, j])); + assert_eq!(transposed.apply(&[i, j]), identity.apply(&[i, j])); } } } @@ -391,7 +396,7 @@ fn unsqueeze_then_squeeze_gives_the_tensor_back() { let unsqueezed = Rc::new(unsqueeze(original.clone(), 1)); assert_eq!(unsqueezed.hrect().shape(), &[3, 1, 2]); - assert_eq!(unsqueezed.get(&[2, 0, 1]), original.get(&[2, 1])); + assert_eq!(unsqueezed.apply(&[2, 0, 1]), original.apply(&[2, 1])); let squeezed = squeeze(unsqueezed, ()); assert_eq!(squeezed.hrect().shape(), &[3, 2]); @@ -414,12 +419,12 @@ fn unsqueeze_can_add_an_axis_before_the_first_or_after_the_last() { fn squeeze_with_no_dims_drops_every_axis_of_extent_one() { let squeezed = squeeze(Rc::new(ones(vec![1, 3, 1, 2])), ()); assert_eq!(squeezed.hrect().shape(), &[3, 2]); - assert_eq!(squeezed.get(&[2, 1]), 1.0); + assert_eq!(squeezed.apply(&[2, 1]), 1.0); // When every axis has extent one, what is left is a scalar. let scalar = squeeze(Rc::new(ones(vec![1, 1])), ()); assert_eq!(scalar.hrect().ndim(), 0); - assert_eq!(scalar.get(&[]), 1.0); + assert_eq!(scalar.apply(&[]), 1.0); } #[test] @@ -548,10 +553,10 @@ fn stack_adds_a_new_axis_at_any_position() { for (t, part) in parts.iter().enumerate() { for i in 0..3 { for j in 0..2 { - let expected = part.get(&[i, j]); - assert_eq!(first.get(&[t, i, j]), expected); - assert_eq!(middle.get(&[i, t, j]), expected); - assert_eq!(last.get(&[i, j, t]), expected); + let expected = part.apply(&[i, j]); + assert_eq!(first.apply(&[t, i, j]), expected); + assert_eq!(middle.apply(&[i, t, j]), expected); + assert_eq!(last.apply(&[i, j, t]), expected); } } } @@ -577,7 +582,7 @@ fn where_reads_input_where_the_condition_holds_and_other_elsewhere() { // A comparison turns a float tensor into a boolean one, to use as a condition. #[test] fn where_can_take_its_condition_from_a_comparison() { - let above_three = Rc::new(LeftComposition::new(physical(), |x: f32| x > 3.0)); + let above_three = Rc::new(compose(ValueFunction::new(|x: f32| x > 3.0), physical())); let picked = where_(above_three, physical(), Rc::new(zeros(vec![6]))); assert_eq!(collect_1d(&picked, 6), vec![0.0, 0.0, 0.0, 4.0, 5.0, 6.0]); @@ -588,7 +593,10 @@ fn where_can_take_its_condition_from_a_comparison() { #[test] fn where_can_take_its_condition_from_a_tensor_with_no_data() { let m = matrix(2, 2, vec![1.0, 2.0, 3.0, 4.0]); - let diagonal = Rc::new(LeftComposition::new(Rc::new(eye(2)), |x: f32| x != 0.0)); + let diagonal = Rc::new(compose( + ValueFunction::new(|x: f32| x != 0.0), + Rc::new(eye(2)), + )); let picked = where_(diagonal, m.clone(), Rc::new(negate(m))); assert_eq!( @@ -627,19 +635,22 @@ fn sparse_coo_tensor_matches_the_torch_documentation() { // the coordinates given one row per dimension instead of one per value. #[test] fn sparse_coo_tensor_is_the_demo_sparse_tensor() { - let by_hand = Sum::new( - physical(), - Rc::new(COO::new( - HRect::new(vec![6, 5]), - vec![ - vec![0, 0], - vec![0, 3], - vec![4, 2], - vec![5, 4], - vec![1, 3], - vec![2, 0], - ], - )), + let by_hand = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(COO::new( + HRect::new(vec![6, 5]), + vec![ + vec![0, 0], + vec![0, 3], + vec![4, 2], + vec![5, 4], + vec![1, 3], + vec![2, 0], + ], + )), + ), ); let built = sparse_coo_tensor( vec![vec![0, 0, 4, 5, 1, 2], vec![0, 3, 2, 4, 3, 0]], @@ -729,7 +740,7 @@ fn sparse_latticed_tensor_lays_out_an_anti_diagonal() { fn sparse_latticed_tensor_lays_out_a_block_diagonal() { // Three 2 by 2 blocks holding 1 to 12, block after block, row by row. let values: Vec = (1..=12).map(|value| value as f32).collect(); - let blocks = Rc::new(RightComposition::new( + let blocks = Rc::new(compose( Rc::new(Physical::new(values)), Rc::new(Stride::new(HRect::new(vec![3, 2, 2]), vec![4, 2, 1])), )); @@ -776,7 +787,7 @@ fn sparse_latticed_tensor_pads_with_its_offset_and_shape() { // stride, as torch's `expand` does, so it still holds only three values. #[test] fn sparse_latticed_tensor_lays_out_a_circulant_matrix() { - let kernel = Rc::new(RightComposition::new( + let kernel = Rc::new(compose( Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), Rc::new(Stride::new(HRect::new(vec![4, 3]), vec![0, 1])), )); @@ -798,7 +809,7 @@ fn sparse_latticed_tensor_lays_out_a_circulant_matrix() { // by zero strides, and the basis puts the two copies side by side on each row. #[test] fn sparse_latticed_tensor_repeats_a_vector_without_copying_it() { - let expanded = Rc::new(RightComposition::new( + let expanded = Rc::new(compose( Rc::new(Physical::new(vec![1.0, 2.0, 3.0])), Rc::new(Stride::new(HRect::new(vec![4, 2, 3]), vec![0, 0, 1])), )); @@ -855,9 +866,9 @@ fn index_functions_work_on_integer_tensors() { vec![vec![1, 4], vec![2, 5], vec![3, 6]] ); assert_eq!(joined.hrect().shape(), &[4, 3]); - assert_eq!(joined.get(&[3, 2]), 6); + assert_eq!(joined.apply(&[3, 2]), 6); assert_eq!(stacked.hrect().shape(), &[2, 2, 3]); - assert_eq!(stacked.get(&[1, 1, 0]), 4); + assert_eq!(stacked.apply(&[1, 1, 0]), 4); } #[test] @@ -879,8 +890,8 @@ fn sum_and_prod_reduce_integer_and_double_tensors() { assert_eq!(collect_1d(&sum(ints.clone(), 0), 3), vec![5, 7, 9]); assert_eq!(collect_1d(&prod(ints, 1), 2), vec![6, 120]); - assert_eq!(sum(doubles.clone(), 0).get(&[]), 2.75); - assert_eq!(prod(doubles, 0).get(&[]), 0.25); + assert_eq!(sum(doubles.clone(), 0).apply(&[]), 2.75); + assert_eq!(prod(doubles, 0).apply(&[]), 0.25); } #[test] diff --git a/tests/logic.rs b/tests/logic.rs index d661468..9189b94 100644 --- a/tests/logic.rs +++ b/tests/logic.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - Affine, COO, Diagonal, EmptyMapping, HRect, HRectSignature, Inf, LeftComposition, Multiset, - Physical, PowersetMapping, Product, RightComposition, Stride, Sum, Sup, Tensor, + Affine, COO, Diagonal, EmptyMapping, Function, HRect, HRectSignature, Inf, Mapping, Multiset, + Physical, Product, Stride, Sum, Sup, Tensor, ValueFunction, compose, multiset_lift, }; fn physical() -> Rc { @@ -11,7 +11,7 @@ fn physical() -> Rc { fn collect_2d(tensor: &dyn Tensor, rows: usize, cols: usize) -> Vec> { (0..rows) - .map(|i| (0..cols).map(|j| tensor.get(&[i, j])).collect()) + .map(|i| (0..cols).map(|j| tensor.apply(&[i, j])).collect()) .collect() } @@ -33,25 +33,25 @@ fn coo_2d() -> Rc { fn physical_reads_its_backing_storage() { let physical = physical(); for i in 0..6 { - assert_eq!(physical.get(&[i]), (i + 1) as f32); + assert_eq!(physical.apply(&[i]), (i + 1) as f32); } } #[test] #[should_panic] fn physical_rejects_out_of_bounds_indices() { - physical().get(&[6]); + physical().apply(&[6]); } #[test] #[should_panic(expected = "single coordinate")] fn physical_rejects_indices_of_the_wrong_rank() { - physical().get(&[0, 0]); + physical().apply(&[0, 0]); } #[test] fn right_composition_reads_through_a_stride() { - let strided = RightComposition::new( + let strided = compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), ); @@ -64,7 +64,7 @@ fn right_composition_reads_through_a_stride() { #[test] fn a_zero_stride_broadcasts_along_its_axis() { - let repeated = RightComposition::new( + let repeated = compose( physical(), Rc::new(Stride::new(HRect::new(vec![4, 6]), vec![0, 1])), ); @@ -80,12 +80,12 @@ fn a_zero_stride_broadcasts_along_its_axis() { #[test] #[should_panic(expected = "indexed by 2 coordinates")] fn composing_over_a_tensor_of_the_wrong_rank_is_rejected() { - let strided = Rc::new(RightComposition::new( + let strided = Rc::new(compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), )); - RightComposition::new( + compose( strided, Rc::new(Stride::new(HRect::new(vec![2, 3]), vec![1, 1])), ); @@ -93,11 +93,11 @@ fn composing_over_a_tensor_of_the_wrong_rank_is_rejected() { #[test] fn left_composition_transforms_values_pointwise() { - let strided = Rc::new(RightComposition::new( + let strided = Rc::new(compose( physical(), Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), )); - let squared = LeftComposition::new(strided, |x| x * x); + let squared = compose(ValueFunction::new(|x| x * x), strided); assert_eq!( collect_2d(&squared, 3, 2), @@ -113,14 +113,17 @@ fn a_sparse_vector_can_be_read_through_a_mapping() { HRect::new(vec![6]), vec![vec![0], vec![4], vec![5]], )); - let sparse = Rc::new(Sum::new(physical(), coo_vector)); + let sparse = Rc::new(compose( + Sum::new(), + compose(multiset_lift(physical()), coo_vector), + )); - assert_eq!(sparse.get(&[0]), 1.0); - assert_eq!(sparse.get(&[4]), 2.0); - assert_eq!(sparse.get(&[5]), 3.0); - assert_eq!(sparse.get(&[1]), 0.0); + assert_eq!(sparse.apply(&[0]), 1.0); + assert_eq!(sparse.apply(&[4]), 2.0); + assert_eq!(sparse.apply(&[5]), 3.0); + assert_eq!(sparse.apply(&[1]), 0.0); - let repeated = RightComposition::new( + let repeated = compose( sparse.clone(), Rc::new(Stride::new(HRect::new(vec![2, 6]), vec![0, 1])), ); @@ -132,7 +135,7 @@ fn a_sparse_vector_can_be_read_through_a_mapping() { #[test] fn sum_over_a_coo_pattern_scatters_the_stored_values() { - let sparse = Sum::new(physical(), coo_2d()); + let sparse = compose(Sum::new(), compose(multiset_lift(physical()), coo_2d())); assert_eq!( collect_2d(&sparse, 6, 5), @@ -151,15 +154,15 @@ fn sum_over_a_coo_pattern_scatters_the_stored_values() { fn each_reduction_falls_back_to_its_neutral_element() { let off_support = [3, 3]; - let sum = Sum::new(physical(), coo_2d()); - let product = Product::new(physical(), coo_2d()); - let sup = Sup::new(physical(), coo_2d()); - let inf = Inf::new(physical(), coo_2d()); + let sum = compose(Sum::new(), compose(multiset_lift(physical()), coo_2d())); + let product = compose(Product::new(), compose(multiset_lift(physical()), coo_2d())); + let sup = compose(Sup::new(), compose(multiset_lift(physical()), coo_2d())); + let inf = compose(Inf::new(), compose(multiset_lift(physical()), coo_2d())); - assert_eq!(sum.get(&off_support), 0.0); - assert_eq!(product.get(&off_support), 1.0); - assert_eq!(sup.get(&off_support), f32::NEG_INFINITY); - assert_eq!(inf.get(&off_support), f32::INFINITY); + assert_eq!(sum.apply(&off_support), 0.0); + assert_eq!(product.apply(&off_support), 1.0); + assert_eq!(sup.apply(&off_support), f32::NEG_INFINITY); + assert_eq!(inf.apply(&off_support), f32::INFINITY); } // The demo's COO patterns all use distinct coordinates, so every multiset it @@ -176,21 +179,24 @@ fn reductions_accumulate_over_every_matching_position() { let values = || Rc::new(Physical::new(vec![2.0, 3.0, 4.0])); let at = [1, 1]; - let sum = Sum::new(values(), overlapping()); - let product = Product::new(values(), overlapping()); - let sup = Sup::new(values(), overlapping()); - let inf = Inf::new(values(), overlapping()); + let sum = compose(Sum::new(), compose(multiset_lift(values()), overlapping())); + let product = compose( + Product::new(), + compose(multiset_lift(values()), overlapping()), + ); + let sup = compose(Sup::new(), compose(multiset_lift(values()), overlapping())); + let inf = compose(Inf::new(), compose(multiset_lift(values()), overlapping())); - assert_eq!(sum.get(&at), 9.0); - assert_eq!(product.get(&at), 24.0); - assert_eq!(sup.get(&at), 4.0); - assert_eq!(inf.get(&at), 2.0); + assert_eq!(sum.apply(&at), 9.0); + assert_eq!(product.apply(&at), 24.0); + assert_eq!(sup.apply(&at), 4.0); + assert_eq!(inf.apply(&at), 2.0); } #[test] fn a_partial_coo_entry_matches_every_index_sharing_its_prefix() { let row_pattern = Rc::new(COO::new(HRect::new(vec![4, 3]), vec![vec![2]])); - let sum = Sum::new(physical(), row_pattern); + let sum = compose(Sum::new(), compose(multiset_lift(physical()), row_pattern)); assert_eq!( collect_2d(&sum, 4, 3), @@ -205,27 +211,36 @@ fn a_partial_coo_entry_matches_every_index_sharing_its_prefix() { #[test] fn diagonal_reduction_builds_a_diagonal_matrix() { - let diag = Sum::new(physical(), Rc::new(Diagonal::new(HRect::new(vec![6, 6])))); + let diag = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(Diagonal::new(HRect::new(vec![6, 6]))), + ), + ); for i in 0..6 { for j in 0..6 { let expected = if i == j { (i + 1) as f32 } else { 0.0 }; - assert_eq!(diag.get(&[i, j]), expected); + assert_eq!(diag.apply(&[i, j]), expected); } } } #[test] fn diagonal_reduction_generalises_to_rank_three() { - let diag = Sum::new( - physical(), - Rc::new(Diagonal::new(HRect::new(vec![6, 6, 6]))), + let diag = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(Diagonal::new(HRect::new(vec![6, 6, 6]))), + ), ); - assert_eq!(diag.get(&[2, 2, 2]), 3.0); - assert_eq!(diag.get(&[2, 2, 1]), 0.0); - assert_eq!(diag.get(&[0, 1, 2]), 0.0); - assert_eq!(diag.get(&[5, 5, 5]), 6.0); + assert_eq!(diag.apply(&[2, 2, 2]), 3.0); + assert_eq!(diag.apply(&[2, 2, 1]), 0.0); + assert_eq!(diag.apply(&[0, 1, 2]), 0.0); + assert_eq!(diag.apply(&[5, 5, 5]), 6.0); } // A product over the empty set is 1 everywhere, so a constant tensor of ones @@ -233,38 +248,50 @@ fn diagonal_reduction_generalises_to_rank_three() { #[test] fn empty_mapping_under_a_product_yields_ones_from_an_empty_store() { let empty = Rc::new(Physical::::new(vec![])); - let ones = Rc::new(Product::new( - empty, - Rc::new(EmptyMapping::new(HRect::new(vec![6]), HRect::new(vec![0]))), + let ones = Rc::new(compose( + Product::new(), + compose( + multiset_lift(empty), + Rc::new(EmptyMapping::new(HRect::new(vec![6]), HRect::new(vec![0]))), + ), )); for i in 0..6 { - assert_eq!(ones.get(&[i]), 1.0); + assert_eq!(ones.apply(&[i]), 1.0); } - let identity = Sum::new(ones, Rc::new(Diagonal::new(HRect::new(vec![6, 6])))); + let identity = compose( + Sum::new(), + compose( + multiset_lift(ones), + Rc::new(Diagonal::new(HRect::new(vec![6, 6]))), + ), + ); for i in 0..6 { for j in 0..6 { let expected = if i == j { 1.0 } else { 0.0 }; - assert_eq!(identity.get(&[i, j]), expected); + assert_eq!(identity.apply(&[i, j]), expected); } } } #[test] fn compositions_nest_to_arbitrary_depth() { - let sparse = Rc::new(Sum::new( - physical(), - Rc::new(COO::new( - HRect::new(vec![13]), - vec![vec![0], vec![4], vec![5], vec![1], vec![10], vec![12]], - )), + let sparse = Rc::new(compose( + Sum::new(), + compose( + multiset_lift(physical()), + Rc::new(COO::new( + HRect::new(vec![13]), + vec![vec![0], vec![4], vec![5], vec![1], vec![10], vec![12]], + )), + ), )); - let repeated = Rc::new(RightComposition::new( + let repeated = Rc::new(compose( sparse, Rc::new(Stride::new(HRect::new(vec![10, 13]), vec![0, 1])), )); - let squared = LeftComposition::new(repeated, |x| x * x); + let squared = compose(ValueFunction::new(|x| x * x), repeated); let expected: Vec = [ 1.0, 4.0, 0.0, 0.0, 2.0, 3.0, 0.0, 0.0, 0.0, 0.0, 5.0, 0.0, 6.0, @@ -279,7 +306,7 @@ fn compositions_nest_to_arbitrary_depth() { } #[test] -fn a_powerset_mapping_can_be_used_directly() { +fn an_index_function_can_be_used_directly() { let diagonal = Diagonal::new(HRect::new(vec![6, 6])); let selected = diagonal.map(&[4, 4]); @@ -293,11 +320,11 @@ fn a_powerset_mapping_can_be_used_directly() { #[test] fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { let values: Vec = (1..=9).map(|value| value as f32).collect(); - let matrix = Rc::new(RightComposition::new( + let matrix = Rc::new(compose( Rc::new(Physical::new(values)), Rc::new(Stride::new(HRect::new(vec![3, 3]), vec![3, 1])), )); - let diagonal = RightComposition::new( + let diagonal = compose( matrix, Rc::new(Affine::new( HRect::new(vec![3]), @@ -307,7 +334,7 @@ fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { )), ); - let read: Vec = (0..3).map(|i| diagonal.get(&[i])).collect(); + let read: Vec = (0..3).map(|i| diagonal.apply(&[i])).collect(); assert_eq!(read, vec![1.0, 5.0, 9.0]); } @@ -316,8 +343,8 @@ struct Twice { signature: HRectSignature, } -// Named in full, as importing it would make `map` ambiguous in the tests above. -impl cten::mapping::HRectPowersetMapping for Twice { +impl Mapping for Twice { + type Output = Multiset>; type Codomain = HRect; fn map(&self, index: &[usize]) -> Multiset> { @@ -336,12 +363,15 @@ fn a_reduction_counts_an_index_as_often_as_it_is_repeated() { let twice = Rc::new(Twice { signature: HRectSignature::new(HRect::new(vec![6]), HRect::new(vec![6])), }); - let doubled = Sum::new(physical(), twice.clone()); - let squared = Product::new(physical(), twice); + let doubled = compose( + Sum::new(), + compose(multiset_lift(physical()), twice.clone()), + ); + let squared = compose(Product::new(), compose(multiset_lift(physical()), twice)); for i in 0..6 { let value = (i + 1) as f32; - assert_eq!(doubled.get(&[i]), 2.0 * value); - assert_eq!(squared.get(&[i]), value * value); + assert_eq!(doubled.apply(&[i]), 2.0 * value); + assert_eq!(squared.apply(&[i]), value * value); } } diff --git a/tests/mapping.rs b/tests/mapping.rs index ab81889..45a642d 100644 --- a/tests/mapping.rs +++ b/tests/mapping.rs @@ -1,9 +1,9 @@ use std::{assert_eq, collections::HashMap, rc::Rc}; use cten::{ - Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexMapping, InvAffine, InvProjection, - JaggedDomain, Multiset, Permutation, Physical, PowersetMapping, Projection, Squeeze, Stride, - Unsqueeze, Where, + Affine, COO, Concat, Diagonal, Domain, EmptyMapping, HRect, InvAffine, InvProjection, + JaggedDomain, Mapping, Multiset, Permutation, Physical, Projection, Squeeze, Stride, Unsqueeze, + Where, }; /// The indices of a multiset in sorted order, to compare multisets whose @@ -18,16 +18,16 @@ fn sorted(indices: Multiset>) -> Vec> { fn a_stride_goes_from_its_hrect_to_the_flat_positions_it_reaches() { let stride = Stride::new(HRect::new(vec![3, 2]), vec![1, 3]); - assert!(stride.domain().contains(&[2, 1])); - assert!(!stride.domain().contains(&[3, 0])); + assert!(stride.signature().domain().contains(&[2, 1])); + assert!(!stride.signature().domain().contains(&[3, 0])); // The largest index is 2 * 1 + 1 * 3 = 5, so six positions are reached. - assert_eq!(stride.codomain().ndim(), 1); - assert!(stride.codomain().contains(&[5])); - assert!(!stride.codomain().contains(&[6])); + assert_eq!(stride.signature().codomain().ndim(), 1); + assert!(stride.signature().codomain().contains(&[5])); + assert!(!stride.signature().codomain().contains(&[6])); // A broadcasting stride reaches fewer positions than its domain has indices. let repeat = Stride::new(HRect::new(vec![4, 6]), vec![0, 1]); - assert!(repeat.codomain().contains(&[5])); - assert!(!repeat.codomain().contains(&[6])); + assert!(repeat.signature().codomain().contains(&[5])); + assert!(!repeat.signature().codomain().contains(&[6])); } #[test] @@ -37,31 +37,31 @@ fn a_coo_goes_from_its_hrect_to_the_positions_it_stores() { vec![vec![0, 0], vec![4, 2], vec![5, 4]], ); - assert!(coo.domain().contains(&[5, 4])); - assert!(!coo.domain().contains(&[6, 0])); - assert!(coo.codomain().contains(&[2])); - assert!(!coo.codomain().contains(&[3])); + assert!(coo.signature().domain().contains(&[5, 4])); + assert!(!coo.signature().domain().contains(&[6, 0])); + assert!(coo.signature().codomain().contains(&[2])); + assert!(!coo.signature().codomain().contains(&[3])); } #[test] fn a_diagonal_reaches_as_far_as_its_shortest_dimension() { let diagonal = Diagonal::new(HRect::new(vec![6, 4, 6])); - assert_eq!(diagonal.domain().ndim(), 3); - assert!(diagonal.domain().contains(&[3, 3, 3])); - assert!(!diagonal.domain().contains(&[3, 5, 3])); - assert!(diagonal.codomain().contains(&[3])); - assert!(!diagonal.codomain().contains(&[4])); + assert_eq!(diagonal.signature().domain().ndim(), 3); + assert!(diagonal.signature().domain().contains(&[3, 3, 3])); + assert!(!diagonal.signature().domain().contains(&[3, 5, 3])); + assert!(diagonal.signature().codomain().contains(&[3])); + assert!(!diagonal.signature().codomain().contains(&[4])); } #[test] fn a_projection_removes_dim_from_index() { let proj = Projection::new(HRect::new(vec![2, 3, 4]), 1); - assert_eq!(proj.codomain().ndim(), 2); - assert!(proj.codomain().contains(&[1, 3])); - assert!(!proj.codomain().contains(&[1, 4])); - assert!(proj.domain().contains(&[1, 2, 3])); + assert_eq!(proj.signature().codomain().ndim(), 2); + assert!(proj.signature().codomain().contains(&[1, 3])); + assert!(!proj.signature().codomain().contains(&[1, 4])); + assert!(proj.signature().domain().contains(&[1, 2, 3])); assert_eq!(proj.map(&[1, 2, 3]), vec![1, 3]) } @@ -69,10 +69,10 @@ fn a_projection_removes_dim_from_index() { fn an_inv_projection_sends_each_index_to_the_line_along_its_dim() { let inv_proj = InvProjection::new(HRect::new(vec![2, 3, 4]), 1); - assert_eq!(inv_proj.domain().ndim(), 2); - assert!(inv_proj.domain().contains(&[1, 3])); - assert!(!inv_proj.domain().contains(&[1, 4])); - assert!(inv_proj.codomain().contains(&[1, 2, 3])); + assert_eq!(inv_proj.signature().domain().ndim(), 2); + assert!(inv_proj.signature().domain().contains(&[1, 3])); + assert!(!inv_proj.signature().domain().contains(&[1, 4])); + assert!(inv_proj.signature().codomain().contains(&[1, 2, 3])); assert_eq!( sorted(inv_proj.map(&[1, 3])), vec![vec![1, 0, 3], vec![1, 1, 3], vec![1, 2, 3]] @@ -90,9 +90,9 @@ fn an_affine_sends_each_index_to_its_lattice_point() { HRect::new(vec![3, 3]), ); - assert!(affine.domain().contains(&[2])); - assert!(!affine.domain().contains(&[3])); - assert!(affine.codomain().contains(&[2, 2])); + assert!(affine.signature().domain().contains(&[2])); + assert!(!affine.signature().domain().contains(&[3])); + assert!(affine.signature().codomain().contains(&[2, 2])); assert_eq!(affine.map(&[0]), vec![0, 2]); assert_eq!(affine.map(&[1]), vec![1, 1]); assert_eq!(affine.map(&[2]), vec![2, 0]); @@ -132,9 +132,9 @@ fn an_inv_affine_sends_each_index_to_the_lattice_points_on_it() { HRect::new(vec![9]), )); - assert!(inv_affine.domain().contains(&[8])); - assert!(!inv_affine.domain().contains(&[9])); - assert!(inv_affine.codomain().contains(&[2])); + assert!(inv_affine.signature().domain().contains(&[8])); + assert!(!inv_affine.signature().domain().contains(&[9])); + assert!(inv_affine.signature().codomain().contains(&[2])); assert_eq!(sorted(inv_affine.map(&[4])), vec![vec![1]]); assert!(inv_affine.map(&[5]).is_empty()); } @@ -321,11 +321,11 @@ fn a_permutation_reorders_its_domain_and_keeps_its_codomain() { // Axis 2 of the codomain becomes axis 0 of the domain, and so on, so the // domain is the codomain's shape read through `axes`. assert_eq!(permutation.axes(), &[2, 0, 1]); - assert_eq!(permutation.domain().ndim(), 3); - assert!(permutation.domain().contains(&[3, 1, 2])); - assert!(!permutation.domain().contains(&[4, 1, 2])); - assert!(permutation.codomain().contains(&[1, 2, 3])); - assert!(!permutation.codomain().contains(&[1, 3, 2])); + assert_eq!(permutation.signature().domain().ndim(), 3); + assert!(permutation.signature().domain().contains(&[3, 1, 2])); + assert!(!permutation.signature().domain().contains(&[4, 1, 2])); + assert!(permutation.signature().codomain().contains(&[1, 2, 3])); + assert!(!permutation.signature().codomain().contains(&[1, 3, 2])); assert_eq!(permutation.map(&[3, 1, 2]), vec![1, 2, 3]); } @@ -333,9 +333,9 @@ fn a_permutation_reorders_its_domain_and_keeps_its_codomain() { fn an_empty_mapping_keeps_the_domain_and_codomain_it_was_given() { let empty = EmptyMapping::new(HRect::new(vec![6]), HRect::new(vec![0])); - assert!(empty.domain().contains(&[5])); - assert!(!empty.domain().contains(&[6])); - assert!(!empty.codomain().contains(&[0])); + assert!(empty.signature().domain().contains(&[5])); + assert!(!empty.signature().domain().contains(&[6])); + assert!(!empty.signature().codomain().contains(&[0])); assert!(empty.map(&[5]).is_empty()); } @@ -343,9 +343,9 @@ fn an_empty_mapping_keeps_the_domain_and_codomain_it_was_given() { fn an_unsqueeze_adds_an_axis_of_extent_one_to_its_codomain() { let unsqueeze = Unsqueeze::new(HRect::new(vec![3, 2]), 1); - assert!(unsqueeze.domain().contains(&[2, 0, 1])); - assert!(!unsqueeze.domain().contains(&[2, 1, 1])); - assert!(unsqueeze.codomain().contains(&[2, 1])); + assert!(unsqueeze.signature().domain().contains(&[2, 0, 1])); + assert!(!unsqueeze.signature().domain().contains(&[2, 1, 1])); + assert!(unsqueeze.signature().codomain().contains(&[2, 1])); assert_eq!(unsqueeze.map(&[2, 0, 1]), vec![2, 1]); } @@ -353,10 +353,10 @@ fn an_unsqueeze_adds_an_axis_of_extent_one_to_its_codomain() { fn a_squeeze_drops_axes_of_extent_one_from_its_codomain() { let squeeze = Squeeze::new(HRect::new(vec![1, 3, 1]), vec![2, 0]); - assert_eq!(squeeze.domain().ndim(), 1); - assert!(squeeze.domain().contains(&[2])); - assert!(!squeeze.domain().contains(&[3])); - assert!(squeeze.codomain().contains(&[0, 2, 0])); + assert_eq!(squeeze.signature().domain().ndim(), 1); + assert!(squeeze.signature().domain().contains(&[2])); + assert!(!squeeze.signature().domain().contains(&[3])); + assert!(squeeze.signature().codomain().contains(&[0, 2, 0])); // The axes were given out of order; each coordinate still lands back in place. assert_eq!(squeeze.map(&[2]), vec![0, 2, 0]); } @@ -375,14 +375,14 @@ fn a_concat_goes_from_a_hrect_to_the_jagged_domain_of_its_parts() { ); // Extents 3 and 1 along axis 1 add up to 4. - assert!(concat.domain().contains(&[1, 3])); - assert!(!concat.domain().contains(&[1, 4])); - assert!(!concat.domain().contains(&[2, 0])); - assert_eq!(concat.codomain().ndim(), 3); + assert!(concat.signature().domain().contains(&[1, 3])); + assert!(!concat.signature().domain().contains(&[1, 4])); + assert!(!concat.signature().domain().contains(&[2, 0])); + assert_eq!(concat.signature().codomain().ndim(), 3); // The part comes first, then the index within it. assert_eq!(concat.map(&[1, 2]), vec![0, 1, 2]); assert_eq!(concat.map(&[1, 3]), vec![1, 1, 0]); - assert!(concat.codomain().contains(&[1, 1, 0])); + assert!(concat.signature().codomain().contains(&[1, 1, 0])); } #[test] @@ -401,8 +401,8 @@ fn a_where_sends_each_index_to_the_part_its_condition_picks() { Rc::new(Physical::new(vec![true, false, true])), ); - assert!(picks.domain().contains(&[2])); - assert!(!picks.domain().contains(&[3])); + assert!(picks.signature().domain().contains(&[2])); + assert!(!picks.signature().domain().contains(&[3])); assert_eq!(picks.map(&[0]), vec![0, 0]); assert_eq!(picks.map(&[1]), vec![1, 1]); assert_eq!(picks.map(&[2]), vec![0, 2]); diff --git a/tests/multiset.rs b/tests/multiset.rs index b96dd42..f2438c1 100644 --- a/tests/multiset.rs +++ b/tests/multiset.rs @@ -1,4 +1,4 @@ -use cten::{InfReduction, Multiset, ProductReduction, Reduction, SumReduction, SupReduction}; +use cten::{Function, Inf, Multiset, Product, Sum, Sup}; #[test] fn a_multiset_prints_with_double_braces() { @@ -18,25 +18,22 @@ fn a_multiset_keeps_every_repeat() { // The example of issue #18: the sum counts the 3 twice. #[test] fn a_reduction_reduces_a_multiset_to_one_value() { - assert_eq!( - SumReduction::reduce(Multiset::new(vec![1.0, 3.0, 3.0])), - 7.0 - ); - assert_eq!(SumReduction::reduce(Multiset::new(vec![1, 3, 3])), 7); - assert_eq!(ProductReduction::reduce(Multiset::new(vec![2, 3, 3])), 18); - assert_eq!(SupReduction::reduce(Multiset::new(vec![2, 7, 5])), 7); - assert_eq!(InfReduction::reduce(Multiset::new(vec![2, 7, 5])), 2); + assert_eq!(Sum::new().apply(&Multiset::new(vec![1.0, 3.0, 3.0])), 7.0); + assert_eq!(Sum::new().apply(&Multiset::new(vec![1, 3, 3])), 7); + assert_eq!(Product::new().apply(&Multiset::new(vec![2, 3, 3])), 18); + assert_eq!(Sup::new().apply(&Multiset::new(vec![2, 7, 5])), 7); + assert_eq!(Inf::new().apply(&Multiset::new(vec![2, 7, 5])), 2); } // Reducing the empty multiset gives the neutral element, and reducing a single // element gives that element. #[test] fn a_reduction_of_no_value_or_one_value_is_as_expected() { - assert_eq!(SumReduction::reduce(Multiset::::default()), 0.0); - assert_eq!(ProductReduction::reduce(Multiset::::default()), 1.0); + assert_eq!(Sum::new().apply(&Multiset::::default()), 0.0); + assert_eq!(Product::new().apply(&Multiset::::default()), 1.0); assert_eq!( - SupReduction::reduce(Multiset::::default()), + Sup::new().apply(&Multiset::::default()), f32::NEG_INFINITY ); - assert_eq!(InfReduction::reduce(Multiset::new(vec![4.0])), 4.0); + assert_eq!(Inf::new().apply(&Multiset::new(vec![4.0])), 4.0); } diff --git a/tests/tensor.rs b/tests/tensor.rs index 56de005..09a7a2f 100644 --- a/tests/tensor.rs +++ b/tests/tensor.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - Bar, BarComposition, BoxedTensor, COO, Diagonal, HRect, InvProjection, JaggedTensor, - LeftComposition, Multiset, Physical, RightComposition, Stride, Sum, Tensor, + Boxed, COO, Compose, Diagonal, Domain, Function, HRect, InvProjection, JaggedTensor, Multiset, + Permutation, Physical, Stride, Sum, ValueFunction, compose, multiset_lift, }; fn physical() -> Rc { @@ -13,92 +13,116 @@ fn physical() -> Rc { fn a_physical_tensor_is_defined_on_the_hrect_of_its_storage() { let physical = physical(); - assert_eq!(physical.domain().ndim(), 1); - assert!(physical.domain().contains(&[5])); - assert!(!physical.domain().contains(&[6])); + assert_eq!(physical.hrect().ndim(), 1); + assert!(physical.hrect().contains(&[5])); + assert!(!physical.hrect().contains(&[6])); assert_eq!(physical.hrect().shape(), &[6]); } #[test] -fn a_right_composition_takes_its_domain_from_its_mapping() { - let strided = Rc::new(RightComposition::new( +fn a_composition_takes_its_domain_from_the_function_applied_first() { + let strided = Rc::new(compose( physical(), - Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), + Stride::new(HRect::new(vec![3, 2]), vec![1, 3]), )); - let transposed = RightComposition::new( - physical(), - Rc::new(Stride::new(HRect::new(vec![2, 3]), vec![3, 1])), - ); + let transposed = compose(physical(), Stride::new(HRect::new(vec![2, 3]), vec![3, 1])); assert_eq!(strided.hrect().shape(), &[3, 2]); - assert!(strided.domain().contains(&[2, 1])); - assert!(!strided.domain().contains(&[2, 2])); + assert!(strided.hrect().contains(&[2, 1])); + assert!(!strided.hrect().contains(&[2, 2])); // The same storage under a mapping of another shape is another domain. assert_eq!(transposed.hrect().shape(), &[2, 3]); - assert!(transposed.domain().contains(&[1, 2])); + assert!(transposed.hrect().contains(&[1, 2])); } #[test] -fn a_reduction_takes_its_domain_from_its_mapping() { - let sparse = Sum::new( - physical(), - Rc::new(COO::new(HRect::new(vec![2, 3]), vec![vec![0, 0]])), +fn a_reduction_tensor_takes_its_domain_from_its_index_function() { + let sparse = compose( + Sum::new(), + compose( + multiset_lift(physical()), + COO::new(HRect::new(vec![2, 3]), vec![vec![0, 0]]), + ), ); - let diagonal = Sum::new( - physical(), - Rc::new(Diagonal::new(HRect::new(vec![4, 4, 4]))), + let diagonal = compose( + Sum::new(), + compose( + multiset_lift(physical()), + Diagonal::new(HRect::new(vec![4, 4, 4])), + ), ); assert_eq!(sparse.hrect().shape(), &[2, 3]); - assert!(!sparse.domain().contains(&[2, 0])); + assert!(!sparse.hrect().contains(&[2, 0])); assert_eq!(diagonal.hrect().shape(), &[4, 4, 4]); } #[test] -fn a_left_composition_takes_its_domain_from_the_tensor_it_transforms() { - let strided = Rc::new(RightComposition::new( +fn a_function_on_values_keeps_the_domain_of_the_tensor_it_follows() { + let strided = Rc::new(compose( physical(), - Rc::new(Stride::new(HRect::new(vec![3, 2]), vec![1, 3])), + Stride::new(HRect::new(vec![3, 2]), vec![1, 3]), )); - let squared = LeftComposition::new(strided.clone(), |x| x * x); + let squared = compose(ValueFunction::new(|x: f32| x * x), strided.clone()); assert_eq!(squared.hrect().shape(), strided.hrect().shape()); - assert!(squared.domain().contains(&[2, 1])); - assert!(!squared.domain().contains(&[3, 0])); - assert_eq!(squared.get(&[2, 1]), 36.0); + assert!(squared.hrect().contains(&[2, 1])); + assert!(!squared.hrect().contains(&[3, 0])); + assert_eq!(squared.apply(&[2, 1]), 36.0); +} + +// Composition is not only for tensors: permuting the axes of indices one way and +// then back is the identity on indices. +#[test] +fn index_maps_compose_like_any_other_functions() { + let forth = Permutation::new(HRect::new(vec![4, 2, 3]), vec![1, 2, 0]); + let back = Permutation::new(HRect::new(vec![2, 3, 4]), vec![2, 0, 1]); + let round_trip = compose(back, forth); + + assert_eq!(round_trip.hrect().shape(), &[2, 3, 4]); + for index in [[0, 0, 0], [1, 2, 3], [0, 1, 2]] { + assert_eq!(round_trip.apply(&index), index); + } } -// A multiset of indices reads a multiset of values, repeats included. +// The multiset extension of a tensor reads a multiset of values at a multiset of +// indices, repeats included. #[test] -fn a_bar_reads_a_multiset_of_values_at_a_multiset_of_indices() { - let bar = Bar::new(physical()); - let read = bar.get(&Multiset::new(vec![vec![0], vec![2], vec![2]])); +fn a_lifted_tensor_reads_a_multiset_of_values_at_a_multiset_of_indices() { + let lifted = multiset_lift(physical()); + let read = lifted.apply(&Multiset::new(vec![vec![0], vec![2], vec![2]])); assert_eq!(read.into_iter().collect::>(), vec![1.0, 3.0, 3.0]); - assert!(bar.get(&Multiset::default()).is_empty()); + assert!(lifted.apply(&Multiset::default()).is_empty()); +} + +#[test] +fn a_lifted_function_on_values_applies_to_each_element() { + let squared = multiset_lift(ValueFunction::new(|x: f32| x * x)); + let read = squared.apply(&Multiset::new(vec![1.0, 2.0, 2.0])); + + assert_eq!(read.into_iter().collect::>(), vec![1.0, 4.0, 4.0]); } /// The matrix of the paper's row sums example, `[[1, 2, 3], [0, 4, 1], [2, 2, 5]]`. -fn figure_matrix() -> Rc> { - Rc::new(RightComposition::new( - Rc::new(Physical::new(vec![ - 1.0, 2.0, 3.0, 0.0, 4.0, 1.0, 2.0, 2.0, 5.0, - ])), - Rc::new(Stride::new(HRect::new(vec![3, 3]), vec![3, 1])), +fn figure_matrix() -> Rc> { + Rc::new(compose( + Physical::new(vec![1.0, 2.0, 3.0, 0.0, 4.0, 1.0, 2.0, 2.0, 5.0]), + Stride::new(HRect::new(vec![3, 3]), vec![3, 1]), )) } // Before reduction, each row index reads the multiset of its row; reducing those // multisets with a sum gives the row sums. #[test] -fn a_bar_composition_is_a_reduction_before_the_reduction() { +fn a_lifted_tensor_after_an_index_function_is_a_reduction_before_the_reduction() { let rows = Rc::new(InvProjection::new(HRect::new(vec![3, 3]), 1)); - let gathered = BarComposition::new(figure_matrix(), rows.clone()); - let row_sums = Sum::new(figure_matrix(), rows); + let gathered = Rc::new(compose(multiset_lift(figure_matrix()), rows)); + let row_sums = compose(Sum::new(), gathered.clone()); assert_eq!(gathered.hrect().shape(), &[3]); let before: Vec> = (0..3) - .map(|i| gathered.get(&[i]).into_iter().collect()) + .map(|i| gathered.apply(&[i]).into_iter().collect()) .collect(); assert_eq!( before, @@ -108,7 +132,7 @@ fn a_bar_composition_is_a_reduction_before_the_reduction() { vec![2.0, 2.0, 5.0] ] ); - let after: Vec = (0..3).map(|i| row_sums.get(&[i])).collect(); + let after: Vec = (0..3).map(|i| row_sums.apply(&[i])).collect(); assert_eq!(after, vec![6.0, 5.0, 9.0]); } @@ -117,35 +141,33 @@ fn a_physical_tensor_holds_values_of_any_type() { let ints = Physical::new(vec![3i64, -1]); let flags = Physical::new(vec![true, false]); - assert_eq!(ints.get(&[1]), -1); - assert!(flags.get(&[0])); - assert!(!flags.get(&[1])); + assert_eq!(ints.apply(&[1]), -1); + assert!(flags.apply(&[0])); + assert!(!flags.apply(&[1])); } -// A transform may change the value type, as a comparison turns numbers into -// booleans. +// A function on values may change their type, as a comparison turns numbers +// into booleans. #[test] -fn a_left_composition_can_change_the_value_type() { - let above_three = LeftComposition::new(physical(), |x: f32| x > 3.0); +fn a_function_on_values_can_change_the_value_type() { + let above_three = compose(ValueFunction::new(|x: f32| x > 3.0), physical()); assert_eq!(above_three.hrect().shape(), &[6]); - assert!(!above_three.get(&[2])); - assert!(above_three.get(&[3])); + assert!(!above_three.apply(&[2])); + assert!(above_three.apply(&[3])); } #[test] fn a_jagged_tensor_reads_each_part_under_its_first_coordinate() { let jagged = JaggedTensor::new(vec![physical(), Rc::new(Physical::new(vec![10.0, 20.0]))]); - assert_eq!(jagged.get(&[0, 4]), 5.0); - assert_eq!(jagged.get(&[1, 1]), 20.0); + assert_eq!(jagged.apply(&[0, 4]), 5.0); + assert_eq!(jagged.apply(&[1, 1]), 20.0); // Its domain is the two storages side by side, each keeping its own length. - assert_eq!(jagged.domain().ndim(), 2); - assert!(jagged.domain().contains(&[0, 5])); - assert!(jagged.domain().contains(&[1, 1])); - assert!(!jagged.domain().contains(&[1, 2])); - assert_eq!( - jagged.jagged_domain().parts(), - &[HRect::new(vec![6]), HRect::new(vec![2])] - ); + let domain = jagged.jagged_domain(); + assert_eq!(domain.ndim(), 2); + assert!(domain.contains(&[0, 5])); + assert!(domain.contains(&[1, 1])); + assert!(!domain.contains(&[1, 2])); + assert_eq!(domain.parts(), &[HRect::new(vec![6]), HRect::new(vec![2])]); } From 29ab710694a60ecefce06e12d44debcda43ec8e6 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Thu, 5 Nov 2026 18:02:20 +0530 Subject: [PATCH 3/4] Make every domain a real set and rename MultisetLift to MultiFunction --- src/display.rs | 35 ++++++----- src/function.rs | 147 ++++++++++++++++++--------------------------- src/functions.rs | 99 +++++++++++++++--------------- src/lib.rs | 12 ++-- src/mapping.rs | 111 +++++++++++++++------------------- src/set.rs | 112 ++++++++++++++++++++++++++++++++++ src/signature.rs | 7 ++- src/tensor.rs | 70 ++++++++++----------- tests/functions.rs | 92 ++++++++++++++-------------- tests/logic.rs | 19 +++--- tests/mapping.rs | 14 ++--- tests/tensor.rs | 70 ++++++++++++++------- 12 files changed, 448 insertions(+), 340 deletions(-) create mode 100644 src/set.rs diff --git a/src/display.rs b/src/display.rs index 1d29377..4a586dc 100644 --- a/src/display.rs +++ b/src/display.rs @@ -1,18 +1,18 @@ -//! Printing of boxed tensors: those whose domain is a hrect. The `Display` -//! implementations live here rather than next to each tensor, since they all -//! render the same way, by walking a shape. +//! Printing of functions on a hrect: tensors, and multiset-valued tensors before +//! a reduction. The `Display` implementations live here rather than next to +//! each tensor, since they all render the same way, by walking a shape. -use std::borrow::Borrow; use std::fmt; use std::fmt::{Display, Write as _}; -use crate::function::{Boxed, Compose, Function}; -use crate::tensor::{BoxedTensor, Physical}; +use crate::domain::HRect; +use crate::function::{Compose, Function}; +use crate::set::Set; +use crate::tensor::{Physical, Tensor}; -/// Renders a boxed tensor. Every `Display` implementation below defers to this, -/// and each is written for boxed tensors only. -pub fn format_boxed(tensor: &dyn BoxedTensor) -> String { - format_tensor(tensor.hrect().shape(), &|index| tensor.apply(index)) +/// Renders a tensor on a hrect. +pub fn format_boxed(tensor: &dyn Tensor) -> String { + format_tensor(tensor.domain().shape(), &|index| tensor.apply(index)) } impl Display for Physical { @@ -21,15 +21,18 @@ impl Display for Physical { } } -// A composition is a boxed tensor when the function applied first is one on a -// hrect and the one applied last gives values that print. -impl> Display for Compose +// A composition prints when the function applied first is one on a hrect and the +// one applied last gives values that print, multisets included. +impl Display for Compose where - G::Output: Borrow, - F::Output: Display, + F: Function, + G: Function, + ::Owned: Display, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(&format_boxed(self)) + f.write_str(&format_tensor(self.domain().shape(), &|index| { + self.apply(index) + })) } } diff --git a/src/function.rs b/src/function.rs index eb19662..5000181 100644 --- a/src/function.rs +++ b/src/function.rs @@ -5,63 +5,42 @@ use std::borrow::Borrow; use std::rc::Rc; -use crate::domain::{Domain, HRect}; -use crate::multiset::Multiset; +use crate::set::{Multisets, Set, Values}; -/// A function from `Input` to `Output`. A tensor is a function from indices to -/// values, an index map one from indices to indices, an index function one from -/// indices to multisets of indices, and a reduction one from multisets of values -/// to values. +/// A function from its domain to its codomain, which may be sets of indices, +/// of values, or of multisets of either. A tensor goes from indices to values, +/// an index map from indices to indices, an index function from indices to +/// multisets of indices, and a reduction from multisets of values to values. pub trait Function { - type Input: ?Sized; - type Output; + type Domain: Set; + type Codomain: Set; - fn apply(&self, x: &Self::Input) -> Self::Output; + fn apply(&self, x: &::Element) -> ::Owned; - /// For a function on indices, or on multisets of indices, the set those - /// indices are drawn from. A function on values has none. - fn domain(&self) -> Option<&dyn Domain> { - None - } - - /// For a function to indices, or to multisets of indices, the set those - /// indices are drawn from. A function to values has none. - fn codomain(&self) -> Option<&dyn Domain> { - None - } -} + fn domain(&self) -> Self::Domain; -/// A function whose domain is a hrect. Only these can be printed, since -/// formatting walks a shape and visits every index of it. -pub trait Boxed: Function { - fn hrect(&self) -> &HRect; + fn codomain(&self) -> Self::Codomain; } // A shared function is the same function, so that one can be used in several // compositions, or have its type erased as an `Rc>`. impl Function for Rc { - type Input = T::Input; - type Output = T::Output; + type Domain = T::Domain; + type Codomain = T::Codomain; - fn apply(&self, x: &Self::Input) -> Self::Output { + fn apply(&self, x: &::Element) -> ::Owned { (**self).apply(x) } - fn domain(&self) -> Option<&dyn Domain> { + fn domain(&self) -> Self::Domain { (**self).domain() } - fn codomain(&self) -> Option<&dyn Domain> { + fn codomain(&self) -> Self::Codomain { (**self).codomain() } } -impl Boxed for Rc { - fn hrect(&self) -> &HRect { - (**self).hrect() - } -} - /// `f ∘ g`: applies `g`, then `f` to the result. pub struct Compose { f: F, @@ -69,88 +48,70 @@ pub struct Compose { } /// Composes `f` after `g`. Composing a tensor after an index map re-indexes it, -/// and a function on values after a tensor applies it element-wise. Where both -/// sides are on indices, the indices `g` lands on must be those `f` takes, which -/// is checked here as far as it can be without enumerating a domain. -pub fn compose(f: F, g: G) -> Compose +/// and a function on values after a tensor applies it element-wise. The +/// codomain of `g` must lie in the domain of `f`. +pub fn compose(f: F, g: G) -> Compose where - G::Output: Borrow, + F: Function, + G: Function, { - if let (Some(codomain), Some(domain)) = (g.codomain(), f.domain()) { - assert_eq!( - codomain.ndim(), - domain.ndim(), - "the inner function lands on {}-dimensional indices, but the outer one is indexed by {} coordinates", - codomain.ndim(), - domain.ndim() - ); - } + let (codomain, domain) = (g.codomain(), f.domain()); + assert!( + codomain.lies_in(&domain), + "the inner function lands in {codomain:?}, which does not lie in the domain {domain:?} of the outer one" + ); Compose { f, g } } -impl Function for Compose +impl Function for Compose where - G::Output: Borrow, + F: Function, + G: Function, { - type Input = G::Input; - type Output = F::Output; + type Domain = G::Domain; + type Codomain = F::Codomain; - fn apply(&self, x: &Self::Input) -> Self::Output { + fn apply(&self, x: &::Element) -> ::Owned { self.f.apply(self.g.apply(x).borrow()) } - fn domain(&self) -> Option<&dyn Domain> { + fn domain(&self) -> Self::Domain { self.g.domain() } - fn codomain(&self) -> Option<&dyn Domain> { + fn codomain(&self) -> Self::Codomain { self.f.codomain() } } -impl Boxed for Compose -where - G::Output: Borrow, -{ - fn hrect(&self) -> &HRect { - self.g.hrect() - } -} - -/// The extension of `f` to multisets, written `f̄` in the paper: it maps a -/// multiset to the multiset of the images of its elements, repeats included. -pub struct MultisetLift { +/// A function from multisets to multisets: the extension of `f` to multisets, +/// written `f̄` in the paper. It maps a multiset to the multiset of the images +/// 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 fn multiset_lift(f: F) -> MultisetLift -where - F::Input: ToOwned, -{ - MultisetLift { f } +pub fn multiset_lift(f: F) -> MultiFunction { + MultiFunction { f } } -// The elements are taken owned, so that a tensor, which takes `[usize]`, lifts to -// take the `Multiset>` an index function returns. -impl Function for MultisetLift -where - F::Input: ToOwned, -{ - type Input = Multiset<::Owned>; - type Output = Multiset; +impl Function for MultiFunction { + type Domain = Multisets; + type Codomain = Multisets; - fn apply(&self, x: &Self::Input) -> Self::Output { + fn apply(&self, x: &::Element) -> ::Owned { x.iter() .map(|element| self.f.apply(element.borrow())) .collect() } - fn domain(&self) -> Option<&dyn Domain> { - self.f.domain() + fn domain(&self) -> Self::Domain { + Multisets::new(self.f.domain()) } - fn codomain(&self) -> Option<&dyn Domain> { - self.f.codomain() + fn codomain(&self) -> Self::Codomain { + Multisets::new(self.f.codomain()) } } @@ -169,10 +130,18 @@ impl ValueFunction { } impl Function for ValueFunction { - type Input = A; - type Output = B; + type Domain = Values; + type Codomain = Values; fn apply(&self, x: &A) -> B { (self.f)(x.clone()) } + + fn domain(&self) -> Values { + Values::new() + } + + fn codomain(&self) -> Values { + Values::new() + } } diff --git a/src/functions.rs b/src/functions.rs index 114b8fc..20dded9 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -8,29 +8,30 @@ use std::ops::{Mul, Neg}; use std::rc::Rc; use crate::domain::{Domain, HRect}; -use crate::function::{Compose, ValueFunction, compose, multiset_lift}; +use crate::function::{Compose, Function, ValueFunction, compose, multiset_lift}; use crate::mapping::{ Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Permutation, Projection, Squeeze, Unsqueeze, Where, }; +use crate::set::Values; use crate::tensor::{ - BoxedTensor, JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, - ReductionTensor, Sum, SumReduction, Tensor, + JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, ReductionTensor, Sum, + SumReduction, Tensor, }; -pub fn square + Copy + 'static, U: Tensor>( +pub fn square + Copy + 'static, U: Function>>( u: U, ) -> Compose, U> { compose(ValueFunction::new(|x: A| x * x), u) } -pub fn negate + Clone + 'static, U: Tensor>( +pub fn negate + Clone + 'static, U: Function>>( u: U, ) -> Compose, U> { compose(ValueFunction::new(|x: A| -x), u) } -pub fn scale + Copy + 'static, U: Tensor>( +pub fn scale + Copy + 'static, U: Function>>( u: U, factor: A, ) -> Compose, U> { @@ -40,7 +41,7 @@ pub fn scale + Copy + 'static, U: Tensor>( /// 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>( +pub fn clamp>>( u: U, low: A, high: A, @@ -60,16 +61,16 @@ pub fn clamp>( /// Reorders the axes of a tensor: `axes[i]` is the axis of `u` that becomes /// axis `i` of the result. -pub fn permute>(u: U, axes: Vec) -> Compose { - let codomain = u.hrect().clone(); +pub fn permute>(u: U, axes: Vec) -> Compose { + let codomain = u.domain(); compose(u, 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>(u: U) -> Compose { - let ndim = u.hrect().ndim(); +pub fn transpose>(u: U) -> Compose { + let ndim = u.domain().ndim(); assert!(ndim >= 2, "transpose needs at least two dimensions"); let mut axes: Vec = (0..ndim).collect(); axes.swap(ndim - 2, ndim - 1); @@ -78,12 +79,12 @@ pub fn transpose>(u: U) -> Compose { /// Moves the axis at `source` to `destination`; the axes left over keep their /// relative order. -pub fn movedim>( +pub fn movedim>( u: U, source: usize, destination: usize, ) -> Compose { - let ndim = u.hrect().ndim(); + let ndim = u.domain().ndim(); assert!(source < ndim, "source axis {source} is out of range"); assert!( destination < ndim, @@ -97,8 +98,8 @@ pub fn movedim>( /// Inserts an axis of extent 1 at `dim`, which can be anywhere from before the /// first axis to after the last. -pub fn unsqueeze>(u: U, dim: usize) -> Compose { - let codomain = u.hrect().clone(); +pub fn unsqueeze>(u: U, dim: usize) -> Compose { + let codomain = u.domain(); compose(u, Unsqueeze::new(codomain, dim)) } @@ -142,8 +143,8 @@ 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>(u: U, dims: impl Dims) -> Compose { - let codomain = u.hrect().clone(); +pub fn squeeze>(u: U, dims: impl Dims) -> Compose { + let codomain = u.domain(); let ndim = codomain.shape().len(); let named = dims.axes(ndim); for (i, &axis) in named.iter().enumerate() { @@ -162,32 +163,32 @@ pub fn squeeze>(u: U, dims: impl Dims) -> Compose( - tensors: Vec>>, + tensors: Vec>>, dim: usize, ) -> Compose, Concat> { assert!(!tensors.is_empty(), "concatenate needs at least one tensor"); let underlying = JaggedTensor::new(tensors); - let codomain = underlying.jagged_domain().clone(); + let codomain = underlying.domain(); compose(underlying, Concat::new(codomain, dim)) } /// 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>>, + tensors: Vec>>, dim: usize, ) -> Compose, Concat> { assert!(!tensors.is_empty(), "stack needs at least one tensor"); - let shape = tensors[0].hrect().clone(); + let shape = tensors[0].domain(); assert!( - tensors.iter().all(|tensor| *tensor.hrect() == shape), + tensors.iter().all(|tensor| tensor.domain() == shape), "stack needs tensors of one shape" ); let unsqueezed = tensors .into_iter() .map(|tensor| { - let unsqueeze = Unsqueeze::new(tensor.hrect().clone(), dim); - Rc::new(compose(tensor, unsqueeze)) as Rc> + let unsqueeze = Unsqueeze::new(tensor.domain(), dim); + Rc::new(compose(tensor, unsqueeze)) as Rc> }) .collect(); concatenate(unsqueezed, dim) @@ -196,28 +197,28 @@ pub fn stack( /// 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_( - condition: Rc>, - input: Rc>, - other: Rc>, + condition: Rc>, + input: Rc>, + other: Rc>, ) -> Compose, Where> { let underlying = JaggedTensor::new(vec![input, other]); - let codomain = underlying.jagged_domain().clone(); + let codomain = underlying.domain(); compose(underlying, Where::new(codomain, condition)) } /// 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>( +pub fn diagonalize>( u: U, rank: usize, ) -> ReductionTensor, U, Diagonal> where SumReduction: Reduction, { - assert_eq!(u.hrect().ndim(), 1, "diagonalize takes a vector"); + 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]; + let extent = u.domain().shape()[0]; let diagonal = Diagonal::new(HRect::new(vec![extent; rank])); compose(Sum::new(), compose(multiset_lift(u), diagonal)) } @@ -266,7 +267,7 @@ pub fn eye( /// `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>( indices: Vec>, values: U, size: Vec, @@ -280,11 +281,11 @@ where "sparse_coo_tensor needs one row of indices per dimension of size" ); assert_eq!( - values.hrect().ndim(), + values.domain().ndim(), 1, "sparse_coo_tensor takes a vector of values" ); - let count = values.hrect().shape()[0]; + let count = values.domain().shape()[0]; for row in &indices { assert_eq!( row.len(), @@ -312,7 +313,7 @@ where /// `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>( +pub fn sparse_latticed_tensor>( physical: U, basis: Vec>, offset: Vec, @@ -321,32 +322,32 @@ pub fn sparse_latticed_tensor>( where SumReduction: Reduction, { - let affine = Affine::new(physical.hrect().clone(), basis, offset, HRect::new(shape)); + let affine = Affine::new(physical.domain(), basis, offset, HRect::new(shape)); compose( Sum::new(), compose(multiset_lift(physical), InvAffine::new(affine)), ) } -pub fn sum>( +pub fn sum>( input: U, dim: usize, ) -> ReductionTensor, U, InvProjection> where SumReduction: Reduction, { - let rows = InvProjection::new(input.hrect().clone(), dim); + let rows = InvProjection::new(input.domain(), dim); compose(Sum::new(), compose(multiset_lift(input), rows)) } -pub fn prod>( +pub fn prod>( input: U, dim: usize, ) -> ReductionTensor, U, InvProjection> where ProductReduction: Reduction, { - let rows = InvProjection::new(input.hrect().clone(), dim); + let rows = InvProjection::new(input.domain(), dim); compose(Product::new(), compose(multiset_lift(input), rows)) } @@ -354,7 +355,7 @@ where /// matrices at each `[m, n, p]`. type Products = ReductionTensor, Compose, Concat>, InvProjection>; -pub fn mm + 'static, V: BoxedTensor + 'static>( +pub fn mm + 'static, V: Tensor + 'static>( mat1: U, mat2: V, ) -> ReductionTensor, Products, InvProjection> @@ -362,17 +363,17 @@ where SumReduction: Reduction, ProductReduction: Reduction, { - 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!(mat1.domain().shape().len(), 2, "mat1 must be a matrix"); + assert_eq!(mat2.domain().shape().len(), 2, "mat2 must be a matrix"); assert_eq!( - mat1.hrect().shape()[1], - mat2.hrect().shape()[0], + mat1.domain().shape()[1], + mat2.domain().shape()[0], "mat1 must have as many columns as mat2 has rows" ); - let m = mat1.hrect().shape()[0]; - let n = mat1.hrect().shape()[1]; - let p = mat2.hrect().shape()[1]; + let m = mat1.domain().shape()[0]; + let n = mat1.domain().shape()[1]; + let p = mat2.domain().shape()[1]; let latent_domain = HRect::new(vec![m, n, p]); let stacked_latent_domain = HRect::new(vec![2, m, n, p]); @@ -381,7 +382,7 @@ where let latent_inv_proj = InvProjection::new(stacked_latent_domain, 0); // l(2, m, n, p) = [m, n, p]^T let psi_inv = InvProjection::new(latent_domain.clone(), 1); // psi(m, n, p) = [m, p]^T - let latent_mats: Vec>> = vec![ + let latent_mats: Vec>> = vec![ Rc::new(compose(mat1, phi1)), // shape [m, n, p] Rc::new(compose(mat2, phi2)), // shape [m, n, p] ]; diff --git a/src/lib.rs b/src/lib.rs index 6ff1837..aa1b6ef 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,24 +8,26 @@ pub mod functions; mod lattice; pub mod mapping; pub mod multiset; +pub mod set; pub mod signature; pub mod tensor; pub use display::format_boxed; pub use domain::{Domain, HRect, JaggedDomain}; -pub use function::{Boxed, Compose, Function, MultisetLift, ValueFunction, compose, multiset_lift}; +pub use function::{Compose, Function, MultiFunction, ValueFunction, compose, multiset_lift}; pub use functions::{ 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::{ - Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Mapping, Permutation, - Projection, Squeeze, Stride, Unsqueeze, Where, + Affine, COO, Concat, Diagonal, EmptyMapping, IndexFunctionSignature, InvAffine, InvProjection, + Mapping, Permutation, Projection, Squeeze, Stride, Unsqueeze, Where, }; pub use multiset::Multiset; +pub use set::{Multisets, Set, Values}; pub use signature::{HRectSignature, Signature}; pub use tensor::{ - BoxedTensor, Inf, InfReduction, JaggedTensor, Physical, Product, ProductReduction, Reduce, - Reduction, ReductionTensor, Sum, SumReduction, Sup, SupReduction, Tensor, + Inf, InfReduction, JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, + ReductionTensor, Sum, SumReduction, Sup, SupReduction, Tensor, }; diff --git a/src/mapping.rs b/src/mapping.rs index 9507fa7..9246309 100644 --- a/src/mapping.rs +++ b/src/mapping.rs @@ -1,46 +1,43 @@ use std::rc::Rc; use crate::domain::{Domain, HRect, JaggedDomain}; -use crate::function::{Boxed, Function}; +use crate::function::Function; use crate::lattice::Preimage; use crate::multiset::Multiset; +use crate::set::{Multisets, Set}; use crate::signature::{HRectSignature, Signature}; -use crate::tensor::BoxedTensor; +use crate::tensor::Tensor; -/// A function on the indices of a hrect: an index map when `Output` is an index -/// of its codomain, and an index function when it is a multiset of them. The -/// codomain is the set those indices are drawn from. Implementing this gets -/// `Function` and `Boxed` for free, with the domain and codomain of the -/// signature. +/// The signature of an index function: from a hrect to the multisets of +/// indices of another. +pub type IndexFunctionSignature = Signature>; + +/// A function on the indices of a hrect: an index map when its codomain is a +/// set of indices, and an index function when it is the multisets of one. +/// Implementing this gets `Function` for free, with the domain and codomain of +/// the signature. pub trait Mapping { - type Output; - type Codomain: Domain; + type Codomain: Set; - fn map(&self, index: &[usize]) -> Self::Output; + fn map(&self, index: &[usize]) -> ::Owned; fn signature(&self) -> &Signature; } impl Function for M { - type Input = [usize]; - type Output = M::Output; + type Domain = HRect; + type Codomain = M::Codomain; - fn apply(&self, index: &[usize]) -> M::Output { + fn apply(&self, index: &[usize]) -> ::Owned { self.map(index) } - fn domain(&self) -> Option<&dyn Domain> { - Some(self.signature().domain()) - } - - fn codomain(&self) -> Option<&dyn Domain> { - Some(self.signature().codomain()) + fn domain(&self) -> HRect { + self.signature().domain().clone() } -} -impl Boxed for M { - fn hrect(&self) -> &HRect { - self.signature().domain() + fn codomain(&self) -> M::Codomain { + self.signature().codomain().clone() } } @@ -81,7 +78,6 @@ fn flat_extent(shape: &[usize], strides: &[usize]) -> usize { } impl Mapping for Stride { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -132,7 +128,6 @@ impl Permutation { } impl Mapping for Permutation { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -172,7 +167,6 @@ impl Unsqueeze { } impl Mapping for Unsqueeze { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -220,7 +214,6 @@ impl Squeeze { } impl Mapping for Squeeze { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -275,7 +268,6 @@ impl Concat { } impl Mapping for Concat { - type Output = Vec; type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { @@ -300,20 +292,20 @@ impl Mapping for Concat { /// tensor's values and not only on the index itself. pub struct Where { signature: Signature, - condition: Rc>, + condition: Rc>, } impl Where { /// Takes the codomain, the two tensors picked from. The domain is their /// common shape, which `condition` must share. - pub fn new(codomain: JaggedDomain, condition: Rc>) -> Self { + pub fn new(codomain: JaggedDomain, condition: Rc>) -> Self { let parts = codomain.parts(); assert_eq!(parts.len(), 2, "where picks from exactly two tensors"); let shape = parts[0].clone(); assert!( - parts[1] == shape && *condition.hrect() == shape, + parts[1] == shape && condition.domain() == shape, "where needs its condition and both tensors to have one shape, got {:?}, {:?} and {:?}", - condition.hrect().shape(), + condition.domain().shape(), parts[0].shape(), parts[1].shape() ); @@ -325,7 +317,6 @@ impl Where { } impl Mapping for Where { - type Output = Vec; type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { @@ -342,32 +333,31 @@ impl Mapping for Where { } pub struct EmptyMapping { - signature: HRectSignature, + signature: IndexFunctionSignature, } impl EmptyMapping { pub fn new(domain: HRect, codomain: HRect) -> Self { Self { - signature: HRectSignature::new(domain, codomain), + signature: IndexFunctionSignature::new(domain, Multisets::new(codomain)), } } } impl Mapping for EmptyMapping { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, _index: &[usize]) -> Multiset> { Multiset::default() } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } pub struct COO { - signature: HRectSignature, + signature: IndexFunctionSignature, inverse_coords: Vec>, } @@ -375,7 +365,7 @@ impl COO { pub fn new(domain: HRect, inverse_coords: Vec>) -> Self { let codomain = HRect::new(vec![inverse_coords.len()]); Self { - signature: HRectSignature::new(domain, codomain), + signature: IndexFunctionSignature::new(domain, Multisets::new(codomain)), inverse_coords, } } @@ -386,8 +376,7 @@ impl COO { } impl Mapping for COO { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { let mut result = Multiset::default(); @@ -403,13 +392,13 @@ impl Mapping for COO { result } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } pub struct Diagonal { - signature: HRectSignature, + signature: IndexFunctionSignature, } impl Diagonal { @@ -418,14 +407,16 @@ impl Diagonal { // is as long as the shortest dimension. let extent = domain.shape().iter().copied().min().unwrap_or(0); Self { - signature: HRectSignature::new(domain, HRect::new(vec![extent])), + signature: IndexFunctionSignature::new( + domain, + Multisets::new(HRect::new(vec![extent])), + ), } } } impl Mapping for Diagonal { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { let mut result = Multiset::default(); @@ -439,7 +430,7 @@ impl Mapping for Diagonal { result } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } @@ -468,7 +459,6 @@ impl Projection { } impl Mapping for Projection { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -485,7 +475,7 @@ impl Mapping for Projection { /// The inverse of the projection dropping axis `dim`. /// e.g. with dim=1, f(m, p) = {[m n p]^T | n in \[N\]} pub struct InvProjection { - signature: HRectSignature, + signature: IndexFunctionSignature, dim: usize, } @@ -500,18 +490,17 @@ impl InvProjection { let mut shape = codomain.shape().to_vec(); shape.remove(dim); Self { - signature: HRectSignature::new(HRect::new(shape), codomain), + signature: IndexFunctionSignature::new(HRect::new(shape), Multisets::new(codomain)), dim, } } } impl Mapping for InvProjection { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { - (0..self.signature.codomain().shape()[self.dim]) + (0..self.signature.codomain().set().shape()[self.dim]) .map(|coord| { let mut result = index.to_vec(); result.insert(self.dim, coord); @@ -520,7 +509,7 @@ impl Mapping for InvProjection { .collect() } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } @@ -580,7 +569,6 @@ impl Affine { } impl Mapping for Affine { - type Output = Vec; type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -612,7 +600,7 @@ impl Mapping for Affine { /// `basis·x = index - offset`, so a read costs about as much as the basis is /// large, however large the tensors. pub struct InvAffine { - signature: HRectSignature, + signature: IndexFunctionSignature, offset: Vec, preimage: Preimage, } @@ -626,9 +614,9 @@ impl InvAffine { } = affine; let preimage = Preimage::new(&basis, signature.domain().shape()); Self { - signature: HRectSignature::new( + signature: IndexFunctionSignature::new( signature.codomain().clone(), - signature.domain().clone(), + Multisets::new(signature.domain().clone()), ), offset, preimage, @@ -637,8 +625,7 @@ impl InvAffine { } impl Mapping for InvAffine { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { let target: Vec = index @@ -649,7 +636,7 @@ impl Mapping for InvAffine { self.preimage.solve(&target).into_iter().collect() } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } diff --git a/src/set.rs b/src/set.rs new file mode 100644 index 0000000..496df33 --- /dev/null +++ b/src/set.rs @@ -0,0 +1,112 @@ +//! The sets functions go between: indices, values, and multisets of either. + +use std::borrow::Borrow; +use std::fmt; +use std::marker::PhantomData; + +use crate::domain::{Domain, HRect, JaggedDomain}; +use crate::multiset::Multiset; + +/// A set that functions go between. A function takes an element of its domain +/// as an `Element`, and returns an element of its codomain as an `Owned`. +pub trait Set: Clone + fmt::Debug { + type Element: ?Sized; + type Owned: Borrow; + + /// Whether every element of `self` lies in `other`. Composing `f` after `g` + /// needs the codomain of `g` to lie in the domain of `f`. + fn lies_in(&self, other: &Self) -> bool; +} + +impl Set for HRect { + type Element = [usize]; + type Owned = Vec; + + fn lies_in(&self, other: &Self) -> bool { + self.ndim() == other.ndim() + && self + .shape() + .iter() + .zip(other.shape()) + .all(|(extent, other)| extent <= other) + } +} + +impl Set for JaggedDomain { + type Element = [usize]; + type Owned = Vec; + + fn lies_in(&self, other: &Self) -> bool { + self.parts().len() <= other.parts().len() + && self + .parts() + .iter() + .zip(other.parts()) + .all(|(part, other)| part.lies_in(other)) + } +} + +/// All the values of type `A`, such as ℝ for `f32`. This is the domain of a +/// function on values, and the codomain of a tensor. +pub struct Values(PhantomData A>); + +impl Values { + pub fn new() -> Self { + Self(PhantomData) + } +} + +// Written out rather than derived, which would require the same of `A`. +impl Default for Values { + fn default() -> Self { + Self::new() + } +} + +impl Clone for Values { + fn clone(&self) -> Self { + Self::new() + } +} + +impl fmt::Debug for Values { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Values<{}>", std::any::type_name::()) + } +} + +impl Set for Values { + type Element = A; + type Owned = A; + + // The type says it all: the compiler has already checked it. + fn lies_in(&self, _other: &Self) -> bool { + true + } +} + +/// The multisets of elements of `S`, written `M(S)` in the paper. +#[derive(Clone, Debug)] +pub struct Multisets { + set: S, +} + +impl Multisets { + pub fn new(set: S) -> Self { + Self { set } + } + + /// The set the elements of the multisets are drawn from. + pub fn set(&self) -> &S { + &self.set + } +} + +impl Set for Multisets { + type Element = Multiset; + type Owned = Multiset; + + fn lies_in(&self, other: &Self) -> bool { + self.set.lies_in(&other.set) + } +} diff --git a/src/signature.rs b/src/signature.rs index ee653cd..29eb3c4 100644 --- a/src/signature.rs +++ b/src/signature.rs @@ -1,14 +1,15 @@ //! The domain and codomain a mapping runs between. -use crate::domain::{Domain, HRect}; +use crate::domain::HRect; +use crate::set::Set; /// The domain and codomain a mapping runs between. -pub struct Signature { +pub struct Signature { domain: D, codomain: C, } -impl Signature { +impl Signature { pub fn new(domain: D, codomain: C) -> Self { Self { domain, codomain } } diff --git a/src/tensor.rs b/src/tensor.rs index 8e16513..e1d9ef7 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,21 +1,17 @@ use std::marker::PhantomData; use std::rc::Rc; -use crate::domain::{Domain, HRect, JaggedDomain}; -use crate::function::{Boxed, Compose, Function, MultisetLift}; +use crate::domain::{HRect, JaggedDomain}; +use crate::function::{Compose, Function, MultiFunction}; use crate::multiset::Multiset; +use crate::set::{Multisets, Set, Values}; -/// A function from indices to values, with values of type `A`, `f32` unless -/// said otherwise. Nothing implements this directly: every function of that -/// type is a tensor. -pub trait Tensor: Function {} +/// A function from the indices of `D` to values of type `A`: by default, a +/// tensor on a hrect, which can be printed. Nothing implements this directly: +/// every function of that kind is a tensor. +pub trait Tensor: Function> {} -impl + ?Sized> Tensor for T {} - -/// A tensor whose domain is a hrect, hence a printable one. -pub trait BoxedTensor: Tensor + Boxed {} - -impl + Boxed + ?Sized> BoxedTensor for T {} +impl> + ?Sized> Tensor for T {} pub struct Physical { values: Vec, @@ -30,8 +26,8 @@ impl Physical { } impl Function for Physical { - type Input = [usize]; - type Output = A; + type Domain = HRect; + type Codomain = Values; fn apply(&self, index: &[usize]) -> A { assert_eq!( @@ -42,14 +38,12 @@ impl Function for Physical { self.values[index[0]].clone() } - fn domain(&self) -> Option<&dyn Domain> { - Some(&self.domain) + fn domain(&self) -> HRect { + self.domain.clone() } -} -impl Boxed for Physical { - fn hrect(&self) -> &HRect { - &self.domain + fn codomain(&self) -> Values { + Values::new() } } @@ -57,31 +51,31 @@ impl Boxed for Physical { /// the paper's stack, which concatenating, stacking and `where` remap into. Its /// domain is jagged rather than a hrect, so it is not printable. pub struct JaggedTensor { - parts: Vec>>, + parts: Vec>>, domain: JaggedDomain, } impl JaggedTensor { - pub fn new(parts: Vec>>) -> Self { - let domain = JaggedDomain::new(parts.iter().map(|part| part.hrect().clone()).collect()); + pub fn new(parts: Vec>>) -> Self { + let domain = JaggedDomain::new(parts.iter().map(|part| part.domain()).collect()); Self { parts, domain } } - - pub fn jagged_domain(&self) -> &JaggedDomain { - &self.domain - } } impl Function for JaggedTensor { - type Input = [usize]; - type Output = A; + type Domain = JaggedDomain; + type Codomain = Values; fn apply(&self, index: &[usize]) -> A { self.parts[index[0]].apply(&index[1..]) } - fn domain(&self) -> Option<&dyn Domain> { - Some(&self.domain) + fn domain(&self) -> JaggedDomain { + self.domain.clone() + } + + fn codomain(&self) -> Values { + Values::new() } } @@ -160,12 +154,20 @@ impl Default for Reduce { } impl, A: Clone> Function for Reduce { - type Input = Multiset; - type Output = A; + type Domain = Multisets>; + type Codomain = Values; fn apply(&self, multiset: &Multiset) -> A { multiset.iter().cloned().fold(R::NEUTRAL, R::combine) } + + fn domain(&self) -> Multisets> { + Multisets::new(Values::new()) + } + + fn codomain(&self) -> Values { + Values::new() + } } pub type Sum = Reduce; @@ -177,4 +179,4 @@ pub type Inf = Reduce; /// underlying tensor `U` and index function `M`: at each index, the reduction /// of the multiset of the values of `U` at the indices `M` maps it to. Built /// as `compose(Sum::new(), compose(multiset_lift(u), psi))`. -pub type ReductionTensor = Compose, M>>; +pub type ReductionTensor = Compose, M>>; diff --git a/tests/functions.rs b/tests/functions.rs index c52f9bd..e3735d4 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -1,10 +1,10 @@ use std::rc::Rc; use cten::{ - Boxed, BoxedTensor, COO, Compose, Domain, Function, HRect, Physical, Stride, Sum, Tensor, - ValueFunction, clamp, compose, concatenate, diagonalize, eye, eye_rect, mm, movedim, - multiset_lift, negate, ones, permute, prod, scale, sparse_coo_tensor, sparse_latticed_tensor, - square, squeeze, stack, sum, transpose, unsqueeze, where_, zeros, + COO, Compose, Domain, Function, HRect, Physical, Stride, Sum, Tensor, ValueFunction, clamp, + compose, concatenate, diagonalize, eye, eye_rect, mm, movedim, multiset_lift, negate, ones, + permute, prod, scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, + transpose, unsqueeze, where_, zeros, }; type Strided = Compose>, Rc>; @@ -142,7 +142,7 @@ fn elementwise_functions_work_on_integer_tensors() { fn an_elementwise_transform_keeps_the_domain() { let squared = square(strided()); - assert_eq!(squared.hrect().shape(), &[3, 2]); + assert_eq!(squared.domain().shape(), &[3, 2]); assert_eq!( collect_2d(&squared, 3, 2), vec![vec![1.0, 16.0], vec![4.0, 25.0], vec![9.0, 36.0]] @@ -155,7 +155,7 @@ fn permute_reorders_axes() { let permuted = permute(cube.clone(), vec![2, 0, 1]); // Axis 2 becomes axis 0, axis 0 becomes axis 1, axis 1 becomes axis 2. - assert_eq!(permuted.hrect().shape(), &[4, 2, 3]); + assert_eq!(permuted.domain().shape(), &[4, 2, 3]); for a in 0..4 { for b in 0..2 { @@ -170,7 +170,7 @@ fn permute_reorders_axes() { fn transpose_swaps_the_last_two_axes() { let transposed = transpose(strided()); - assert_eq!(transposed.hrect().shape(), &[2, 3]); + assert_eq!(transposed.domain().shape(), &[2, 3]); assert_eq!( collect_2d(&transposed, 2, 3), vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]] @@ -182,7 +182,7 @@ fn transpose_leaves_leading_axes_alone() { let cube = cube(); let transposed = transpose(cube.clone()); - assert_eq!(transposed.hrect().shape(), &[2, 4, 3]); + assert_eq!(transposed.domain().shape(), &[2, 4, 3]); for i in 0..2 { for j in 0..4 { @@ -198,7 +198,7 @@ fn transpose_is_an_involution() { let original = strided(); let restored = transpose(Rc::new(transpose(original.clone()))); - assert_eq!(restored.hrect().shape(), &[3, 2]); + assert_eq!(restored.domain().shape(), &[3, 2]); assert_eq!(collect_2d(&restored, 3, 2), collect_2d(&*original, 3, 2)); } @@ -218,7 +218,7 @@ fn transpose_applies_to_a_tensor_with_no_strides() { )); let transposed = transpose(sparse.clone()); - assert_eq!(transposed.hrect().shape(), &[3, 2]); + assert_eq!(transposed.domain().shape(), &[3, 2]); assert_eq!(transposed.apply(&[1, 0]), 1.0); assert_eq!(transposed.apply(&[2, 1]), 2.0); assert_eq!(transposed.apply(&[0, 0]), 0.0); @@ -236,7 +236,7 @@ fn movedim_follows_the_pytorch_convention() { let moved = movedim(cube.clone(), 0, 2); // Moving axis 0 to the end leaves the others in order: (2,3,4) -> (3,4,2). - assert_eq!(moved.hrect().shape(), &[3, 4, 2]); + assert_eq!(moved.domain().shape(), &[3, 4, 2]); for a in 0..3 { for b in 0..4 { @@ -252,7 +252,7 @@ fn movedim_is_reversible() { let cube = cube(); let restored = movedim(Rc::new(movedim(cube.clone(), 0, 2)), 2, 0); - assert_eq!(restored.hrect().shape(), &[2, 3, 4]); + assert_eq!(restored.domain().shape(), &[2, 3, 4]); for i in 0..2 { for j in 0..3 { for k in 0..4 { @@ -266,7 +266,7 @@ fn movedim_is_reversible() { fn diagonalize_builds_a_diagonal_matrix() { let diagonal = diagonalize(physical(), 2); - assert_eq!(diagonal.hrect().shape(), &[6, 6]); + assert_eq!(diagonal.domain().shape(), &[6, 6]); for i in 0..6 { for j in 0..6 { let expected = if i == j { (i + 1) as f32 } else { 0.0 }; @@ -279,7 +279,7 @@ fn diagonalize_builds_a_diagonal_matrix() { fn diagonalize_generalises_to_higher_rank() { let diagonal = diagonalize(physical(), 3); - assert_eq!(diagonal.hrect().shape(), &[6, 6, 6]); + assert_eq!(diagonal.domain().shape(), &[6, 6, 6]); assert_eq!(diagonal.apply(&[2, 2, 2]), 3.0); assert_eq!(diagonal.apply(&[2, 2, 1]), 0.0); assert_eq!(diagonal.apply(&[0, 1, 2]), 0.0); @@ -293,7 +293,7 @@ fn zeros_and_ones_are_constant_at_any_rank() { let flat_zeros = zeros(vec![2, 3]); let flat_ones = ones(vec![2, 3]); - assert_eq!(flat_zeros.hrect().shape(), &[2, 3]); + assert_eq!(flat_zeros.domain().shape(), &[2, 3]); assert_eq!(collect_2d(&flat_zeros, 2, 3), vec![vec![0.0; 3]; 2]); assert_eq!(collect_2d(&flat_ones, 2, 3), vec![vec![1.0; 3]; 2]); @@ -307,7 +307,7 @@ fn zeros_and_ones_are_constant_at_any_rank() { fn a_rank_zero_constant_is_a_scalar() { let one = ones(vec![]); - assert_eq!(one.hrect().ndim(), 0); + assert_eq!(one.domain().ndim(), 0); assert_eq!(one.apply(&[]), 1.0); assert_eq!(zeros(vec![]).apply(&[]), 0.0); } @@ -316,7 +316,7 @@ fn a_rank_zero_constant_is_a_scalar() { fn eye_is_the_identity() { let identity = eye(6); - assert_eq!(identity.hrect().shape(), &[6, 6]); + assert_eq!(identity.domain().shape(), &[6, 6]); for i in 0..6 { for j in 0..6 { let expected = if i == j { 1.0 } else { 0.0 }; @@ -330,13 +330,13 @@ fn eye_rect_puts_ones_on_the_diagonal_of_a_rectangle() { let wide = eye_rect(2, 4); let tall = eye_rect(4, 2); - assert_eq!(wide.hrect().shape(), &[2, 4]); + assert_eq!(wide.domain().shape(), &[2, 4]); assert_eq!( collect_2d(&wide, 2, 4), vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]] ); - assert_eq!(tall.hrect().shape(), &[4, 2]); + assert_eq!(tall.domain().shape(), &[4, 2]); assert_eq!( collect_2d(&tall, 4, 2), vec![ @@ -395,18 +395,18 @@ fn unsqueeze_then_squeeze_gives_the_tensor_back() { let original = strided(); let unsqueezed = Rc::new(unsqueeze(original.clone(), 1)); - assert_eq!(unsqueezed.hrect().shape(), &[3, 1, 2]); + assert_eq!(unsqueezed.domain().shape(), &[3, 1, 2]); assert_eq!(unsqueezed.apply(&[2, 0, 1]), original.apply(&[2, 1])); let squeezed = squeeze(unsqueezed, ()); - assert_eq!(squeezed.hrect().shape(), &[3, 2]); + assert_eq!(squeezed.domain().shape(), &[3, 2]); assert_eq!(collect_2d(&squeezed, 3, 2), collect_2d(&*original, 3, 2)); } #[test] fn unsqueeze_can_add_an_axis_before_the_first_or_after_the_last() { - assert_eq!(unsqueeze(strided(), 0).hrect().shape(), &[1, 3, 2]); - assert_eq!(unsqueeze(strided(), 2).hrect().shape(), &[3, 2, 1]); + assert_eq!(unsqueeze(strided(), 0).domain().shape(), &[1, 3, 2]); + assert_eq!(unsqueeze(strided(), 2).domain().shape(), &[3, 2, 1]); // A vector with an axis in front is a matrix of one row. let row = unsqueeze(physical(), 0); assert_eq!( @@ -418,12 +418,12 @@ fn unsqueeze_can_add_an_axis_before_the_first_or_after_the_last() { #[test] fn squeeze_with_no_dims_drops_every_axis_of_extent_one() { let squeezed = squeeze(Rc::new(ones(vec![1, 3, 1, 2])), ()); - assert_eq!(squeezed.hrect().shape(), &[3, 2]); + assert_eq!(squeezed.domain().shape(), &[3, 2]); assert_eq!(squeezed.apply(&[2, 1]), 1.0); // When every axis has extent one, what is left is a scalar. let scalar = squeeze(Rc::new(ones(vec![1, 1])), ()); - assert_eq!(scalar.hrect().ndim(), 0); + assert_eq!(scalar.domain().ndim(), 0); assert_eq!(scalar.apply(&[]), 1.0); } @@ -432,10 +432,10 @@ fn squeeze_drops_only_the_axes_it_is_given() { let tensor = || Rc::new(ones(vec![1, 3, 1, 2])); let named: &[usize] = &[0, 2]; - assert_eq!(squeeze(tensor(), 2).hrect().shape(), &[1, 3, 2]); - assert_eq!(squeeze(tensor(), [0, 2]).hrect().shape(), &[3, 2]); - assert_eq!(squeeze(tensor(), vec![2, 0]).hrect().shape(), &[3, 2]); - assert_eq!(squeeze(tensor(), named).hrect().shape(), &[3, 2]); + assert_eq!(squeeze(tensor(), 2).domain().shape(), &[1, 3, 2]); + assert_eq!(squeeze(tensor(), [0, 2]).domain().shape(), &[3, 2]); + assert_eq!(squeeze(tensor(), vec![2, 0]).domain().shape(), &[3, 2]); + assert_eq!(squeeze(tensor(), named).domain().shape(), &[3, 2]); } // torch leaves a named axis that is longer than one as it is. @@ -443,10 +443,10 @@ fn squeeze_drops_only_the_axes_it_is_given() { fn squeeze_leaves_a_named_axis_that_is_not_of_extent_one_alone() { let unsqueezed = Rc::new(unsqueeze(strided(), 0)); - assert_eq!(squeeze(unsqueezed.clone(), 0).hrect().shape(), &[3, 2]); - assert_eq!(squeeze(unsqueezed.clone(), 1).hrect().shape(), &[1, 3, 2]); + assert_eq!(squeeze(unsqueezed.clone(), 0).domain().shape(), &[3, 2]); + assert_eq!(squeeze(unsqueezed.clone(), 1).domain().shape(), &[1, 3, 2]); // Of several named axes, only those of extent one go. - assert_eq!(squeeze(unsqueezed, [0, 1]).hrect().shape(), &[3, 2]); + assert_eq!(squeeze(unsqueezed, [0, 1]).domain().shape(), &[3, 2]); } #[test] @@ -472,7 +472,7 @@ fn squeeze_rejects_an_axis_named_twice() { fn concatenate_joins_tensors_of_different_lengths_end_to_end() { let joined = concatenate(vec![physical(), Rc::new(Physical::new(vec![7.0, 8.0]))], 0); - assert_eq!(joined.hrect().shape(), &[8]); + assert_eq!(joined.domain().shape(), &[8]); assert_eq!( collect_1d(&joined, 8), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] @@ -496,7 +496,7 @@ fn concatenate_joins_matrices_along_either_axis() { let rows = concatenate(vec![strided(), strided()], 0); let columns = concatenate(vec![strided(), strided()], 1); - assert_eq!(rows.hrect().shape(), &[6, 2]); + assert_eq!(rows.domain().shape(), &[6, 2]); assert_eq!( collect_2d(&rows, 6, 2), vec![ @@ -508,7 +508,7 @@ fn concatenate_joins_matrices_along_either_axis() { vec![3.0, 6.0], ] ); - assert_eq!(columns.hrect().shape(), &[3, 4]); + assert_eq!(columns.domain().shape(), &[3, 4]); assert_eq!( collect_2d(&columns, 3, 4), vec![ @@ -524,7 +524,7 @@ fn concatenate_passes_over_a_part_of_length_zero() { let empty = || Rc::new(Physical::new(vec![])); let joined = concatenate(vec![empty(), physical(), empty()], 0); - assert_eq!(joined.hrect().shape(), &[6]); + assert_eq!(joined.domain().shape(), &[6]); assert_eq!(collect_1d(&joined, 6), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); } @@ -542,14 +542,14 @@ fn concatenate_rejects_an_empty_list() { #[test] fn stack_adds_a_new_axis_at_any_position() { - let parts: Vec> = vec![strided(), Rc::new(negate(strided()))]; + let parts: Vec> = vec![strided(), Rc::new(negate(strided()))]; let first = stack(parts.clone(), 0); let middle = stack(parts.clone(), 1); let last = stack(parts.clone(), 2); - assert_eq!(first.hrect().shape(), &[2, 3, 2]); - assert_eq!(middle.hrect().shape(), &[3, 2, 2]); - assert_eq!(last.hrect().shape(), &[3, 2, 2]); + assert_eq!(first.domain().shape(), &[2, 3, 2]); + assert_eq!(middle.domain().shape(), &[3, 2, 2]); + assert_eq!(last.domain().shape(), &[3, 2, 2]); for (t, part) in parts.iter().enumerate() { for i in 0..3 { for j in 0..2 { @@ -624,7 +624,7 @@ fn sparse_coo_tensor_matches_the_torch_documentation() { vec![2, 4], ); - assert_eq!(sparse.hrect().shape(), &[2, 4]); + assert_eq!(sparse.domain().shape(), &[2, 4]); assert_eq!( collect_2d(&sparse, 2, 4), vec![vec![0.0, 0.0, 3.0, 0.0], vec![4.0, 0.0, 5.0, 0.0]] @@ -658,7 +658,7 @@ fn sparse_coo_tensor_is_the_demo_sparse_tensor() { vec![6, 5], ); - assert_eq!(built.hrect().shape(), &[6, 5]); + assert_eq!(built.domain().shape(), &[6, 5]); assert_eq!(collect_2d(&built, 6, 5), collect_2d(&by_hand, 6, 5)); } @@ -706,7 +706,7 @@ fn sparse_latticed_tensor_lays_out_a_diagonal() { vec![3, 3], ); - assert_eq!(diagonal.hrect().shape(), &[3, 3]); + assert_eq!(diagonal.domain().shape(), &[3, 3]); assert_eq!( collect_2d(&diagonal, 3, 3), vec![ @@ -847,7 +847,7 @@ fn mm_multiplies_matrices() { let mat2 = matrix(3, 2, vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]); let product = mm(mat1, mat2); - assert_eq!(product.hrect().shape(), &[2, 2]); + assert_eq!(product.domain().shape(), &[2, 2]); assert_eq!( collect_2d(&product, 2, 2), vec![vec![58.0, 64.0], vec![139.0, 154.0]] @@ -865,9 +865,9 @@ fn index_functions_work_on_integer_tensors() { collect_2d(&transpose(m), 3, 2), vec![vec![1, 4], vec![2, 5], vec![3, 6]] ); - assert_eq!(joined.hrect().shape(), &[4, 3]); + assert_eq!(joined.domain().shape(), &[4, 3]); assert_eq!(joined.apply(&[3, 2]), 6); - assert_eq!(stacked.hrect().shape(), &[2, 2, 3]); + assert_eq!(stacked.domain().shape(), &[2, 2, 3]); assert_eq!(stacked.apply(&[1, 1, 0]), 4); } diff --git a/tests/logic.rs b/tests/logic.rs index 9189b94..ed1c956 100644 --- a/tests/logic.rs +++ b/tests/logic.rs @@ -1,8 +1,9 @@ use std::rc::Rc; use cten::{ - Affine, COO, Diagonal, EmptyMapping, Function, HRect, HRectSignature, Inf, Mapping, Multiset, - Physical, Product, Stride, Sum, Sup, Tensor, ValueFunction, compose, multiset_lift, + Affine, COO, Diagonal, EmptyMapping, Function, HRect, IndexFunctionSignature, Inf, Mapping, + Multiset, Multisets, Physical, Product, Stride, Sum, Sup, Tensor, ValueFunction, compose, + multiset_lift, }; fn physical() -> Rc { @@ -78,7 +79,7 @@ fn a_zero_stride_broadcasts_along_its_axis() { // A mapping into flat positions needs a one-dimensional tensor under it. That // used to be a matter of index types; it is now checked when composing. #[test] -#[should_panic(expected = "indexed by 2 coordinates")] +#[should_panic(expected = "does not lie in")] fn composing_over_a_tensor_of_the_wrong_rank_is_rejected() { let strided = Rc::new(compose( physical(), @@ -340,18 +341,17 @@ fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { /// An index function sending each index to the multiset holding it twice. struct Twice { - signature: HRectSignature, + signature: IndexFunctionSignature, } impl Mapping for Twice { - type Output = Multiset>; - type Codomain = HRect; + type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { Multiset::new(vec![index.to_vec(), index.to_vec()]) } - fn signature(&self) -> &HRectSignature { + fn signature(&self) -> &IndexFunctionSignature { &self.signature } } @@ -361,7 +361,10 @@ impl Mapping for Twice { #[test] fn a_reduction_counts_an_index_as_often_as_it_is_repeated() { let twice = Rc::new(Twice { - signature: HRectSignature::new(HRect::new(vec![6]), HRect::new(vec![6])), + signature: IndexFunctionSignature::new( + HRect::new(vec![6]), + Multisets::new(HRect::new(vec![6])), + ), }); let doubled = compose( Sum::new(), diff --git a/tests/mapping.rs b/tests/mapping.rs index 45a642d..d7b5106 100644 --- a/tests/mapping.rs +++ b/tests/mapping.rs @@ -39,8 +39,8 @@ fn a_coo_goes_from_its_hrect_to_the_positions_it_stores() { assert!(coo.signature().domain().contains(&[5, 4])); assert!(!coo.signature().domain().contains(&[6, 0])); - assert!(coo.signature().codomain().contains(&[2])); - assert!(!coo.signature().codomain().contains(&[3])); + assert!(coo.signature().codomain().set().contains(&[2])); + assert!(!coo.signature().codomain().set().contains(&[3])); } #[test] @@ -50,8 +50,8 @@ fn a_diagonal_reaches_as_far_as_its_shortest_dimension() { assert_eq!(diagonal.signature().domain().ndim(), 3); assert!(diagonal.signature().domain().contains(&[3, 3, 3])); assert!(!diagonal.signature().domain().contains(&[3, 5, 3])); - assert!(diagonal.signature().codomain().contains(&[3])); - assert!(!diagonal.signature().codomain().contains(&[4])); + assert!(diagonal.signature().codomain().set().contains(&[3])); + assert!(!diagonal.signature().codomain().set().contains(&[4])); } #[test] @@ -72,7 +72,7 @@ fn an_inv_projection_sends_each_index_to_the_line_along_its_dim() { assert_eq!(inv_proj.signature().domain().ndim(), 2); assert!(inv_proj.signature().domain().contains(&[1, 3])); assert!(!inv_proj.signature().domain().contains(&[1, 4])); - assert!(inv_proj.signature().codomain().contains(&[1, 2, 3])); + assert!(inv_proj.signature().codomain().set().contains(&[1, 2, 3])); assert_eq!( sorted(inv_proj.map(&[1, 3])), vec![vec![1, 0, 3], vec![1, 1, 3], vec![1, 2, 3]] @@ -134,7 +134,7 @@ fn an_inv_affine_sends_each_index_to_the_lattice_points_on_it() { assert!(inv_affine.signature().domain().contains(&[8])); assert!(!inv_affine.signature().domain().contains(&[9])); - assert!(inv_affine.signature().codomain().contains(&[2])); + assert!(inv_affine.signature().codomain().set().contains(&[2])); assert_eq!(sorted(inv_affine.map(&[4])), vec![vec![1]]); assert!(inv_affine.map(&[5]).is_empty()); } @@ -335,7 +335,7 @@ fn an_empty_mapping_keeps_the_domain_and_codomain_it_was_given() { assert!(empty.signature().domain().contains(&[5])); assert!(!empty.signature().domain().contains(&[6])); - assert!(!empty.signature().codomain().contains(&[0])); + assert!(!empty.signature().codomain().set().contains(&[0])); assert!(empty.map(&[5]).is_empty()); } diff --git a/tests/tensor.rs b/tests/tensor.rs index 09a7a2f..f9f5ede 100644 --- a/tests/tensor.rs +++ b/tests/tensor.rs @@ -1,8 +1,8 @@ use std::rc::Rc; use cten::{ - Boxed, COO, Compose, Diagonal, Domain, Function, HRect, InvProjection, JaggedTensor, Multiset, - Permutation, Physical, Stride, Sum, ValueFunction, compose, multiset_lift, + COO, Compose, Diagonal, Domain, Function, HRect, InvProjection, JaggedTensor, Multiset, + Multisets, Permutation, Physical, Stride, Sum, ValueFunction, Values, compose, multiset_lift, }; fn physical() -> Rc { @@ -13,10 +13,10 @@ fn physical() -> Rc { fn a_physical_tensor_is_defined_on_the_hrect_of_its_storage() { let physical = physical(); - assert_eq!(physical.hrect().ndim(), 1); - assert!(physical.hrect().contains(&[5])); - assert!(!physical.hrect().contains(&[6])); - assert_eq!(physical.hrect().shape(), &[6]); + assert_eq!(physical.domain().ndim(), 1); + assert!(physical.domain().contains(&[5])); + assert!(!physical.domain().contains(&[6])); + assert_eq!(physical.domain().shape(), &[6]); } #[test] @@ -27,12 +27,12 @@ fn a_composition_takes_its_domain_from_the_function_applied_first() { )); let transposed = compose(physical(), Stride::new(HRect::new(vec![2, 3]), vec![3, 1])); - assert_eq!(strided.hrect().shape(), &[3, 2]); - assert!(strided.hrect().contains(&[2, 1])); - assert!(!strided.hrect().contains(&[2, 2])); + assert_eq!(strided.domain().shape(), &[3, 2]); + assert!(strided.domain().contains(&[2, 1])); + assert!(!strided.domain().contains(&[2, 2])); // The same storage under a mapping of another shape is another domain. - assert_eq!(transposed.hrect().shape(), &[2, 3]); - assert!(transposed.hrect().contains(&[1, 2])); + assert_eq!(transposed.domain().shape(), &[2, 3]); + assert!(transposed.domain().contains(&[1, 2])); } #[test] @@ -52,9 +52,9 @@ fn a_reduction_tensor_takes_its_domain_from_its_index_function() { ), ); - assert_eq!(sparse.hrect().shape(), &[2, 3]); - assert!(!sparse.hrect().contains(&[2, 0])); - assert_eq!(diagonal.hrect().shape(), &[4, 4, 4]); + assert_eq!(sparse.domain().shape(), &[2, 3]); + assert!(!sparse.domain().contains(&[2, 0])); + assert_eq!(diagonal.domain().shape(), &[4, 4, 4]); } #[test] @@ -65,9 +65,9 @@ fn a_function_on_values_keeps_the_domain_of_the_tensor_it_follows() { )); let squared = compose(ValueFunction::new(|x: f32| x * x), strided.clone()); - assert_eq!(squared.hrect().shape(), strided.hrect().shape()); - assert!(squared.hrect().contains(&[2, 1])); - assert!(!squared.hrect().contains(&[3, 0])); + assert_eq!(squared.domain().shape(), strided.domain().shape()); + assert!(squared.domain().contains(&[2, 1])); + assert!(!squared.domain().contains(&[3, 0])); assert_eq!(squared.apply(&[2, 1]), 36.0); } @@ -79,7 +79,7 @@ fn index_maps_compose_like_any_other_functions() { let back = Permutation::new(HRect::new(vec![2, 3, 4]), vec![2, 0, 1]); let round_trip = compose(back, forth); - assert_eq!(round_trip.hrect().shape(), &[2, 3, 4]); + assert_eq!(round_trip.domain().shape(), &[2, 3, 4]); for index in [[0, 0, 0], [1, 2, 3], [0, 1, 2]] { assert_eq!(round_trip.apply(&index), index); } @@ -120,7 +120,7 @@ fn a_lifted_tensor_after_an_index_function_is_a_reduction_before_the_reduction() let gathered = Rc::new(compose(multiset_lift(figure_matrix()), rows)); let row_sums = compose(Sum::new(), gathered.clone()); - assert_eq!(gathered.hrect().shape(), &[3]); + assert_eq!(gathered.domain().shape(), &[3]); let before: Vec> = (0..3) .map(|i| gathered.apply(&[i]).into_iter().collect()) .collect(); @@ -152,7 +152,7 @@ fn a_physical_tensor_holds_values_of_any_type() { fn a_function_on_values_can_change_the_value_type() { let above_three = compose(ValueFunction::new(|x: f32| x > 3.0), physical()); - assert_eq!(above_three.hrect().shape(), &[6]); + assert_eq!(above_three.domain().shape(), &[6]); assert!(!above_three.apply(&[2])); assert!(above_three.apply(&[3])); } @@ -164,10 +164,38 @@ fn a_jagged_tensor_reads_each_part_under_its_first_coordinate() { assert_eq!(jagged.apply(&[0, 4]), 5.0); assert_eq!(jagged.apply(&[1, 1]), 20.0); // Its domain is the two storages side by side, each keeping its own length. - let domain = jagged.jagged_domain(); + let domain = jagged.domain(); assert_eq!(domain.ndim(), 2); assert!(domain.contains(&[0, 5])); assert!(domain.contains(&[1, 1])); assert!(!domain.contains(&[1, 2])); assert_eq!(domain.parts(), &[HRect::new(vec![6]), HRect::new(vec![2])]); } + +// Every function has a domain and a codomain, functions on values and on +// multisets included: a lifted tensor goes from the multisets of its indices to +// multisets of values, and an index function to multisets of indices. +#[test] +fn every_function_has_a_domain_and_a_codomain() { + let lifted = multiset_lift(physical()); + let lifted_domain: Multisets = lifted.domain(); + let _: Multisets> = lifted.codomain(); + assert_eq!(lifted_domain.set(), &HRect::new(vec![6])); + + let diagonal = Diagonal::new(HRect::new(vec![6, 6])); + let diagonal_codomain: Multisets = diagonal.codomain(); + assert_eq!(diagonal_codomain.set(), &HRect::new(vec![6])); + + let _: Multisets> = Sum::new().domain(); + let _: Values = Sum::new().codomain(); + let _: Values = ValueFunction::new(|x: f32| x * x).domain(); +} + +// The codomain of the inner function must lie in the domain of the outer one: +// here the index function reaches seven stored values, and there are six. +#[test] +#[should_panic(expected = "does not lie in")] +fn composing_over_too_few_stored_values_is_rejected() { + let coo = COO::new(HRect::new(vec![7]), (0..7).map(|i| vec![i]).collect()); + compose(multiset_lift(physical()), coo); +} From 3292db3a921f5835ff778ef254c9af1bfa8363a5 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Fri, 6 Nov 2026 08:50:15 +0530 Subject: [PATCH 4/4] Rename types and files as suggested in review --- src/display.rs | 8 +- src/domain.rs | 75 ---------------- src/function.rs | 10 +-- src/functions.rs | 51 ++++++----- src/{mapping.rs => index_functions.rs} | 78 +++++++++-------- src/lib.rs | 26 +++--- src/reductions.rs | 103 ++++++++++++++++++++++ src/{set.rs => sets.rs} | 80 ++++++++++++++++- src/signature.rs | 10 +-- src/tensor.rs | 105 +---------------------- tests/functions.rs | 12 +-- tests/{mapping.rs => index_functions.rs} | 6 +- tests/logic.rs | 18 ++-- tests/{domain.rs => sets.rs} | 2 +- tests/tensor.rs | 8 +- 15 files changed, 299 insertions(+), 293 deletions(-) delete mode 100644 src/domain.rs rename src/{mapping.rs => index_functions.rs} (89%) create mode 100644 src/reductions.rs rename src/{set.rs => sets.rs} (55%) rename tests/{mapping.rs => index_functions.rs} (98%) rename tests/{domain.rs => sets.rs} (97%) diff --git a/src/display.rs b/src/display.rs index 4a586dc..7f38a3e 100644 --- a/src/display.rs +++ b/src/display.rs @@ -5,9 +5,9 @@ use std::fmt; use std::fmt::{Display, Write as _}; -use crate::domain::HRect; -use crate::function::{Compose, Function}; -use crate::set::Set; +use crate::function::{Composition, Function}; +use crate::sets::HRect; +use crate::sets::Set; use crate::tensor::{Physical, Tensor}; /// Renders a tensor on a hrect. @@ -23,7 +23,7 @@ impl Display for Physical { // A composition prints when the function applied first is one on a hrect and the // one applied last gives values that print, multisets included. -impl Display for Compose +impl Display for Composition where F: Function, G: Function, diff --git a/src/domain.rs b/src/domain.rs deleted file mode 100644 index b121832..0000000 --- a/src/domain.rs +++ /dev/null @@ -1,75 +0,0 @@ -/// A set of indices, an index being a vector of coordinates. Tensors are -/// defined on a domain, and mappings go from a domain to a codomain. -pub trait Domain { - /// Number of coordinates the indices of this domain have. - fn ndim(&self) -> usize; - - fn contains(&self, index: &[usize]) -> bool; -} - -/// A hyperrectangle of indices: coordinate `i` runs over `0..shape[i]`. -#[derive(Clone, PartialEq, Eq, Debug)] -pub struct HRect { - shape: Vec, -} - -impl HRect { - pub fn new(shape: Vec) -> Self { - Self { shape } - } - - pub fn shape(&self) -> &[usize] { - &self.shape - } -} - -impl Domain for HRect { - fn ndim(&self) -> usize { - self.shape.len() - } - - fn contains(&self, index: &[usize]) -> bool { - index.len() == self.shape.len() - && index - .iter() - .zip(&self.shape) - .all(|(coord, extent)| coord < extent) - } -} - -/// The disjoint union of several hrects: `[t, rest..]` belongs to it when `rest` -/// is in hrect `t`. This is the index set of a jagged tensor, whose parts may -/// have different shapes. -#[derive(Clone, PartialEq, Eq, Debug)] -pub struct JaggedDomain { - parts: Vec, -} - -impl JaggedDomain { - pub fn new(parts: Vec) -> Self { - assert!(!parts.is_empty(), "a jagged domain needs at least one part"); - let ndim = parts[0].ndim(); - assert!( - parts.iter().all(|part| part.ndim() == ndim), - "the parts of a jagged domain must all have the same number of coordinates" - ); - Self { parts } - } - - pub fn parts(&self) -> &[HRect] { - &self.parts - } -} - -impl Domain for JaggedDomain { - fn ndim(&self) -> usize { - 1 + self.parts[0].ndim() - } - - fn contains(&self, index: &[usize]) -> bool { - match index.split_first() { - Some((&part, rest)) => part < self.parts.len() && self.parts[part].contains(rest), - None => false, - } - } -} diff --git a/src/function.rs b/src/function.rs index 5000181..baf5cb4 100644 --- a/src/function.rs +++ b/src/function.rs @@ -5,7 +5,7 @@ use std::borrow::Borrow; use std::rc::Rc; -use crate::set::{Multisets, Set, Values}; +use crate::sets::{Multisets, Set, Values}; /// A function from its domain to its codomain, which may be sets of indices, /// of values, or of multisets of either. A tensor goes from indices to values, @@ -42,7 +42,7 @@ impl Function for Rc { } /// `f ∘ g`: applies `g`, then `f` to the result. -pub struct Compose { +pub struct Composition { f: F, g: G, } @@ -50,7 +50,7 @@ pub struct Compose { /// Composes `f` after `g`. Composing a tensor after an index map re-indexes it, /// and a function on values after a tensor applies it element-wise. The /// codomain of `g` must lie in the domain of `f`. -pub fn compose(f: F, g: G) -> Compose +pub fn compose(f: F, g: G) -> Composition where F: Function, G: Function, @@ -60,10 +60,10 @@ where codomain.lies_in(&domain), "the inner function lands in {codomain:?}, which does not lie in the domain {domain:?} of the outer one" ); - Compose { f, g } + Composition { f, g } } -impl Function for Compose +impl Function for Composition where F: Function, G: Function, diff --git a/src/functions.rs b/src/functions.rs index 20dded9..14692c8 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -1,40 +1,39 @@ //! Named constructors for the compositions the library can already express. -//! Everything here is assembled from `function`, `mapping` and `tensor`; none -//! of it adds machinery of its own. +//! Everything here is assembled from `function`, `index_functions`, +//! `reductions` and `tensor`; none of it adds machinery of its own. use std::assert_eq; use std::fmt::Display; use std::ops::{Mul, Neg}; use std::rc::Rc; -use crate::domain::{Domain, HRect}; -use crate::function::{Compose, Function, ValueFunction, compose, multiset_lift}; -use crate::mapping::{ +use crate::function::{Composition, Function, ValueFunction, compose, multiset_lift}; +use crate::index_functions::{ Affine, COO, Concat, Diagonal, EmptyMapping, InvAffine, InvProjection, Permutation, Projection, Squeeze, Unsqueeze, Where, }; -use crate::set::Values; -use crate::tensor::{ - JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, ReductionTensor, Sum, - SumReduction, Tensor, +use crate::reductions::{ + Product, ProductReduction, Reduction, ReductionFunction, Sum, SumReduction, }; +use crate::sets::{HRect, IndexSet, Values}; +use crate::tensor::{JaggedTensor, Physical, ReductionTensor, Tensor}; pub fn square + Copy + 'static, U: Function>>( u: U, -) -> Compose, U> { +) -> Composition, U> { compose(ValueFunction::new(|x: A| x * x), u) } pub fn negate + Clone + 'static, U: Function>>( u: U, -) -> Compose, U> { +) -> Composition, U> { compose(ValueFunction::new(|x: A| -x), u) } pub fn scale + Copy + 'static, U: Function>>( u: U, factor: A, -) -> Compose, U> { +) -> Composition, U> { compose(ValueFunction::new(move |x: A| x * factor), u) } @@ -45,7 +44,7 @@ pub fn clamp Compose, U> { +) -> Composition, U> { assert!(low <= high, "clamp needs low <= high, got {low} and {high}"); let clamped = ValueFunction::new(move |x: A| { if x < low { @@ -61,7 +60,7 @@ pub fn clamp>(u: U, axes: Vec) -> Compose { +pub fn permute>(u: U, axes: Vec) -> Composition { let codomain = u.domain(); compose(u, Permutation::new(codomain, axes)) } @@ -69,7 +68,7 @@ pub fn permute>(u: U, axes: Vec) -> Compose>(u: U) -> Compose { +pub fn transpose>(u: U) -> Composition { let ndim = u.domain().ndim(); assert!(ndim >= 2, "transpose needs at least two dimensions"); let mut axes: Vec = (0..ndim).collect(); @@ -83,7 +82,7 @@ pub fn movedim>( u: U, source: usize, destination: usize, -) -> Compose { +) -> Composition { let ndim = u.domain().ndim(); assert!(source < ndim, "source axis {source} is out of range"); assert!( @@ -98,7 +97,7 @@ pub fn movedim>( /// Inserts an axis of extent 1 at `dim`, which can be anywhere from before the /// first axis to after the last. -pub fn unsqueeze>(u: U, dim: usize) -> Compose { +pub fn unsqueeze>(u: U, dim: usize) -> Composition { let codomain = u.domain(); compose(u, Unsqueeze::new(codomain, dim)) } @@ -143,7 +142,7 @@ 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>(u: U, dims: impl Dims) -> Compose { +pub fn squeeze>(u: U, dims: impl Dims) -> Composition { let codomain = u.domain(); let ndim = codomain.shape().len(); let named = dims.axes(ndim); @@ -165,7 +164,7 @@ pub fn squeeze>(u: U, dims: impl Dims) -> Compose { pub fn concatenate( tensors: Vec>>, dim: usize, -) -> Compose, Concat> { +) -> Composition, Concat> { assert!(!tensors.is_empty(), "concatenate needs at least one tensor"); let underlying = JaggedTensor::new(tensors); let codomain = underlying.domain(); @@ -177,7 +176,7 @@ pub fn concatenate( pub fn stack( tensors: Vec>>, dim: usize, -) -> Compose, Concat> { +) -> Composition, Concat> { assert!(!tensors.is_empty(), "stack needs at least one tensor"); let shape = tensors[0].domain(); assert!( @@ -200,7 +199,7 @@ pub fn where_( condition: Rc>, input: Rc>, other: Rc>, -) -> Compose, Where> { +) -> Composition, Where> { let underlying = JaggedTensor::new(vec![input, other]); let codomain = underlying.domain(); compose(underlying, Where::new(codomain, condition)) @@ -227,11 +226,11 @@ where /// index, so a constant tensor needs no underlying data at all. fn neutral>( shape: Vec, -) -> ReductionTensor, Physical, EmptyMapping> { - let mapping = EmptyMapping::new(HRect::new(shape), HRect::new(vec![0])); +) -> ReductionTensor, Physical, EmptyMapping> { + let empty = EmptyMapping::new(HRect::new(shape), HRect::new(vec![0])); compose( - Reduce::new(), - compose(multiset_lift(Physical::new(vec![])), mapping), + ReductionFunction::new(), + compose(multiset_lift(Physical::new(vec![])), empty), ) } @@ -353,7 +352,7 @@ where /// The tensors `mm` multiplies and then sums, stacked: the product of the two /// matrices at each `[m, n, p]`. -type Products = ReductionTensor, Compose, Concat>, InvProjection>; +type Products = ReductionTensor, Composition, Concat>, InvProjection>; pub fn mm + 'static, V: Tensor + 'static>( mat1: U, diff --git a/src/mapping.rs b/src/index_functions.rs similarity index 89% rename from src/mapping.rs rename to src/index_functions.rs index 9246309..273d5ce 100644 --- a/src/mapping.rs +++ b/src/index_functions.rs @@ -1,22 +1,21 @@ use std::rc::Rc; -use crate::domain::{Domain, HRect, JaggedDomain}; use crate::function::Function; use crate::lattice::Preimage; use crate::multiset::Multiset; -use crate::set::{Multisets, Set}; +use crate::sets::{HRect, IndexSet, JaggedDomain, Multisets, Set}; use crate::signature::{HRectSignature, Signature}; use crate::tensor::Tensor; -/// The signature of an index function: from a hrect to the multisets of +/// The signature of a multi-index function: from a hrect to the multisets of /// indices of another. -pub type IndexFunctionSignature = Signature>; +pub type MultiIndexFunctionSignature = Signature>; -/// A function on the indices of a hrect: an index map when its codomain is a -/// set of indices, and an index function when it is the multisets of one. +/// A function on the indices of a hrect, sending each index to an index of its +/// codomain, or, for a multi-index function, to a multiset of them. /// Implementing this gets `Function` for free, with the domain and codomain of /// the signature. -pub trait Mapping { +pub trait IndexFunction { type Codomain: Set; fn map(&self, index: &[usize]) -> ::Owned; @@ -24,7 +23,7 @@ pub trait Mapping { fn signature(&self) -> &Signature; } -impl Function for M { +impl Function for M { type Domain = HRect; type Codomain = M::Codomain; @@ -77,7 +76,7 @@ fn flat_extent(shape: &[usize], strides: &[usize]) -> usize { .sum::() } -impl Mapping for Stride { +impl IndexFunction for Stride { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -100,7 +99,7 @@ pub struct Permutation { impl Permutation { /// Takes the codomain and derives the domain, the other way round from the - /// mappings above: the shape being permuted is what is known, and the + /// index functions above: the shape being permuted is what is known, and the /// permuted shape follows from it. `axes[i]` is the axis of the codomain /// that becomes axis `i` of the domain. pub fn new(codomain: HRect, axes: Vec) -> Self { @@ -127,7 +126,7 @@ impl Permutation { } } -impl Mapping for Permutation { +impl IndexFunction for Permutation { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -166,7 +165,7 @@ impl Unsqueeze { } } -impl Mapping for Unsqueeze { +impl IndexFunction for Unsqueeze { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -213,7 +212,7 @@ impl Squeeze { } } -impl Mapping for Squeeze { +impl IndexFunction for Squeeze { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -267,7 +266,7 @@ impl Concat { } } -impl Mapping for Concat { +impl IndexFunction for Concat { type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { @@ -288,7 +287,7 @@ impl Mapping for Concat { } /// Reads the first of two tensors where `condition` holds and the second where -/// it does not. Unlike the mappings above, where an index is sent depends on a +/// it does not. Unlike the index functions above, where an index is sent depends on a /// tensor's values and not only on the index itself. pub struct Where { signature: Signature, @@ -316,7 +315,7 @@ impl Where { } } -impl Mapping for Where { +impl IndexFunction for Where { type Codomain = JaggedDomain; fn map(&self, index: &[usize]) -> Vec { @@ -333,31 +332,31 @@ impl Mapping for Where { } pub struct EmptyMapping { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, } impl EmptyMapping { pub fn new(domain: HRect, codomain: HRect) -> Self { Self { - signature: IndexFunctionSignature::new(domain, Multisets::new(codomain)), + signature: MultiIndexFunctionSignature::new(domain, Multisets::new(codomain)), } } } -impl Mapping for EmptyMapping { +impl IndexFunction for EmptyMapping { type Codomain = Multisets; fn map(&self, _index: &[usize]) -> Multiset> { Multiset::default() } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } pub struct COO { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, inverse_coords: Vec>, } @@ -365,7 +364,7 @@ impl COO { pub fn new(domain: HRect, inverse_coords: Vec>) -> Self { let codomain = HRect::new(vec![inverse_coords.len()]); Self { - signature: IndexFunctionSignature::new(domain, Multisets::new(codomain)), + signature: MultiIndexFunctionSignature::new(domain, Multisets::new(codomain)), inverse_coords, } } @@ -375,7 +374,7 @@ impl COO { } } -impl Mapping for COO { +impl IndexFunction for COO { type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { @@ -392,13 +391,13 @@ impl Mapping for COO { result } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } pub struct Diagonal { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, } impl Diagonal { @@ -407,7 +406,7 @@ impl Diagonal { // is as long as the shortest dimension. let extent = domain.shape().iter().copied().min().unwrap_or(0); Self { - signature: IndexFunctionSignature::new( + signature: MultiIndexFunctionSignature::new( domain, Multisets::new(HRect::new(vec![extent])), ), @@ -415,7 +414,7 @@ impl Diagonal { } } -impl Mapping for Diagonal { +impl IndexFunction for Diagonal { type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { @@ -430,7 +429,7 @@ impl Mapping for Diagonal { result } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } @@ -458,7 +457,7 @@ impl Projection { } } -impl Mapping for Projection { +impl IndexFunction for Projection { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -475,7 +474,7 @@ impl Mapping for Projection { /// The inverse of the projection dropping axis `dim`. /// e.g. with dim=1, f(m, p) = {[m n p]^T | n in \[N\]} pub struct InvProjection { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, dim: usize, } @@ -490,13 +489,16 @@ impl InvProjection { let mut shape = codomain.shape().to_vec(); shape.remove(dim); Self { - signature: IndexFunctionSignature::new(HRect::new(shape), Multisets::new(codomain)), + signature: MultiIndexFunctionSignature::new( + HRect::new(shape), + Multisets::new(codomain), + ), dim, } } } -impl Mapping for InvProjection { +impl IndexFunction for InvProjection { type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { @@ -509,7 +511,7 @@ impl Mapping for InvProjection { .collect() } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } @@ -568,7 +570,7 @@ impl Affine { } } -impl Mapping for Affine { +impl IndexFunction for Affine { type Codomain = HRect; fn map(&self, index: &[usize]) -> Vec { @@ -600,7 +602,7 @@ impl Mapping for Affine { /// `basis·x = index - offset`, so a read costs about as much as the basis is /// large, however large the tensors. pub struct InvAffine { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, offset: Vec, preimage: Preimage, } @@ -614,7 +616,7 @@ impl InvAffine { } = affine; let preimage = Preimage::new(&basis, signature.domain().shape()); Self { - signature: IndexFunctionSignature::new( + signature: MultiIndexFunctionSignature::new( signature.codomain().clone(), Multisets::new(signature.domain().clone()), ), @@ -624,7 +626,7 @@ impl InvAffine { } } -impl Mapping for InvAffine { +impl IndexFunction for InvAffine { type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { @@ -636,7 +638,7 @@ impl Mapping for InvAffine { self.preimage.solve(&target).into_iter().collect() } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } diff --git a/src/lib.rs b/src/lib.rs index aa1b6ef..89ab9db 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,32 +2,32 @@ //! compositions or functional transforms. pub mod display; -pub mod domain; pub mod function; pub mod functions; +pub mod index_functions; mod lattice; -pub mod mapping; pub mod multiset; -pub mod set; +pub mod reductions; +pub mod sets; pub mod signature; pub mod tensor; pub use display::format_boxed; -pub use domain::{Domain, HRect, JaggedDomain}; -pub use function::{Compose, Function, MultiFunction, ValueFunction, compose, multiset_lift}; +pub use function::{Composition, Function, MultiFunction, ValueFunction, compose, multiset_lift}; pub use functions::{ 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::{ - Affine, COO, Concat, Diagonal, EmptyMapping, IndexFunctionSignature, InvAffine, InvProjection, - Mapping, Permutation, Projection, Squeeze, Stride, Unsqueeze, Where, +pub use index_functions::{ + Affine, COO, Concat, Diagonal, EmptyMapping, IndexFunction, InvAffine, InvProjection, + MultiIndexFunctionSignature, Permutation, Projection, Squeeze, Stride, Unsqueeze, Where, }; pub use multiset::Multiset; -pub use set::{Multisets, Set, Values}; -pub use signature::{HRectSignature, Signature}; -pub use tensor::{ - Inf, InfReduction, JaggedTensor, Physical, Product, ProductReduction, Reduce, Reduction, - ReductionTensor, Sum, SumReduction, Sup, SupReduction, Tensor, +pub use reductions::{ + Inf, InfReduction, Product, ProductReduction, Reduction, ReductionFunction, Sum, SumReduction, + Sup, SupReduction, }; +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/reductions.rs b/src/reductions.rs new file mode 100644 index 0000000..a3d0e94 --- /dev/null +++ b/src/reductions.rs @@ -0,0 +1,103 @@ +//! Reductions: functions from multisets of values to values, `Ω` in the paper. + +use std::marker::PhantomData; + +use crate::function::Function; +use crate::multiset::Multiset; +use crate::sets::{Multisets, Values}; + +/// A way of combining values of type `A` into one, from a neutral element. +pub trait Reduction { + const NEUTRAL: A; + fn combine(a: A, b: A) -> A; +} + +pub struct SumReduction; +pub struct ProductReduction; +pub struct SupReduction; +pub struct InfReduction; + +// The four reductions for each value type they apply to: a sum starts from 0, a +// product from 1, and a sup and an inf from the bottom and top of the type. +macro_rules! reductions { + ($($value:ty: $zero:expr, $one:expr, $bottom:expr, $top:expr;)*) => {$( + impl Reduction<$value> for SumReduction { + const NEUTRAL: $value = $zero; + fn combine(a: $value, b: $value) -> $value { + a + b + } + } + + impl Reduction<$value> for ProductReduction { + const NEUTRAL: $value = $one; + fn combine(a: $value, b: $value) -> $value { + a * b + } + } + + impl Reduction<$value> for SupReduction { + const NEUTRAL: $value = $bottom; + fn combine(a: $value, b: $value) -> $value { + a.max(b) + } + } + + impl Reduction<$value> for InfReduction { + const NEUTRAL: $value = $top; + fn combine(a: $value, b: $value) -> $value { + a.min(b) + } + } + )*}; +} + +reductions! { + f32: 0.0, 1.0, f32::NEG_INFINITY, f32::INFINITY; + f64: 0.0, 1.0, f64::NEG_INFINITY, f64::INFINITY; + i64: 0, 1, i64::MIN, i64::MAX; +} + +/// The reduction `R` as a function from multisets of values to values: the +/// neutral element for the empty multiset, the element of a multiset of one, +/// and otherwise its elements combined in any order, as `combine` is +/// associative and commutative. +pub struct ReductionFunction { + _marker: PhantomData (R, A)>, +} + +impl ReductionFunction { + pub fn new() -> Self { + Self { + _marker: PhantomData, + } + } +} + +// Written out rather than derived, which would require `R: Default` and `A: Default`. +impl Default for ReductionFunction { + fn default() -> Self { + Self::new() + } +} + +impl, A: Clone> Function for ReductionFunction { + type Domain = Multisets>; + type Codomain = Values; + + fn apply(&self, multiset: &Multiset) -> A { + multiset.iter().cloned().fold(R::NEUTRAL, R::combine) + } + + fn domain(&self) -> Multisets> { + Multisets::new(Values::new()) + } + + fn codomain(&self) -> Values { + Values::new() + } +} + +pub type Sum = ReductionFunction; +pub type Product = ReductionFunction; +pub type Sup = ReductionFunction; +pub type Inf = ReductionFunction; diff --git a/src/set.rs b/src/sets.rs similarity index 55% rename from src/set.rs rename to src/sets.rs index 496df33..c0c800c 100644 --- a/src/set.rs +++ b/src/sets.rs @@ -1,10 +1,10 @@ -//! The sets functions go between: indices, values, and multisets of either. +//! The sets functions go between: sets of indices, sets of values, and sets of +//! multisets of either. use std::borrow::Borrow; use std::fmt; use std::marker::PhantomData; -use crate::domain::{Domain, HRect, JaggedDomain}; use crate::multiset::Multiset; /// A set that functions go between. A function takes an element of its domain @@ -18,6 +18,45 @@ pub trait Set: Clone + fmt::Debug { fn lies_in(&self, other: &Self) -> bool; } +/// A set of indices, an index being a vector of coordinates. Tensors go from +/// one, and index functions from one to another or to its multisets. +pub trait IndexSet { + /// Number of coordinates the indices of this set have. + fn ndim(&self) -> usize; + + fn contains(&self, index: &[usize]) -> bool; +} + +/// A hyperrectangle of indices: coordinate `i` runs over `0..shape[i]`. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct HRect { + shape: Vec, +} + +impl HRect { + pub fn new(shape: Vec) -> Self { + Self { shape } + } + + pub fn shape(&self) -> &[usize] { + &self.shape + } +} + +impl IndexSet for HRect { + fn ndim(&self) -> usize { + self.shape.len() + } + + fn contains(&self, index: &[usize]) -> bool { + index.len() == self.shape.len() + && index + .iter() + .zip(&self.shape) + .all(|(coord, extent)| coord < extent) + } +} + impl Set for HRect { type Element = [usize]; type Owned = Vec; @@ -32,6 +71,43 @@ impl Set for HRect { } } +/// The disjoint union of several hrects: `[t, rest..]` belongs to it when `rest` +/// is in hrect `t`. This is the domain of a jagged tensor, whose parts may have +/// different shapes. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct JaggedDomain { + parts: Vec, +} + +impl JaggedDomain { + pub fn new(parts: Vec) -> Self { + assert!(!parts.is_empty(), "a jagged domain needs at least one part"); + let ndim = parts[0].ndim(); + assert!( + parts.iter().all(|part| part.ndim() == ndim), + "the parts of a jagged domain must all have the same number of coordinates" + ); + Self { parts } + } + + pub fn parts(&self) -> &[HRect] { + &self.parts + } +} + +impl IndexSet for JaggedDomain { + fn ndim(&self) -> usize { + 1 + self.parts[0].ndim() + } + + fn contains(&self, index: &[usize]) -> bool { + match index.split_first() { + Some((&part, rest)) => part < self.parts.len() && self.parts[part].contains(rest), + None => false, + } + } +} + impl Set for JaggedDomain { type Element = [usize]; type Owned = Vec; diff --git a/src/signature.rs b/src/signature.rs index 29eb3c4..4b5b29f 100644 --- a/src/signature.rs +++ b/src/signature.rs @@ -1,9 +1,9 @@ -//! The domain and codomain a mapping runs between. +//! The domain and codomain an index function runs between. -use crate::domain::HRect; -use crate::set::Set; +use crate::sets::HRect; +use crate::sets::Set; -/// The domain and codomain a mapping runs between. +/// The domain and codomain an index function runs between. pub struct Signature { domain: D, codomain: C, @@ -23,5 +23,5 @@ impl Signature { } } -/// The signature of a mapping between two hyperrectangles. +/// The signature of an index function between two hyperrectangles. pub type HRectSignature = Signature; diff --git a/src/tensor.rs b/src/tensor.rs index e1d9ef7..03a255d 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,10 +1,7 @@ -use std::marker::PhantomData; use std::rc::Rc; -use crate::domain::{HRect, JaggedDomain}; -use crate::function::{Compose, Function, MultiFunction}; -use crate::multiset::Multiset; -use crate::set::{Multisets, Set, Values}; +use crate::function::{Composition, Function, MultiFunction}; +use crate::sets::{HRect, JaggedDomain, Set, Values}; /// A function from the indices of `D` to values of type `A`: by default, a /// tensor on a hrect, which can be printed. Nothing implements this directly: @@ -79,104 +76,8 @@ impl Function for JaggedTensor { } } -/// A way of combining values of type `A` into one, from a neutral element. -pub trait Reduction { - const NEUTRAL: A; - fn combine(a: A, b: A) -> A; -} - -pub struct SumReduction; -pub struct ProductReduction; -pub struct SupReduction; -pub struct InfReduction; - -// The four reductions for each value type they apply to: a sum starts from 0, a -// product from 1, and a sup and an inf from the bottom and top of the type. -macro_rules! reductions { - ($($value:ty: $zero:expr, $one:expr, $bottom:expr, $top:expr;)*) => {$( - impl Reduction<$value> for SumReduction { - const NEUTRAL: $value = $zero; - fn combine(a: $value, b: $value) -> $value { - a + b - } - } - - impl Reduction<$value> for ProductReduction { - const NEUTRAL: $value = $one; - fn combine(a: $value, b: $value) -> $value { - a * b - } - } - - impl Reduction<$value> for SupReduction { - const NEUTRAL: $value = $bottom; - fn combine(a: $value, b: $value) -> $value { - a.max(b) - } - } - - impl Reduction<$value> for InfReduction { - const NEUTRAL: $value = $top; - fn combine(a: $value, b: $value) -> $value { - a.min(b) - } - } - )*}; -} - -reductions! { - f32: 0.0, 1.0, f32::NEG_INFINITY, f32::INFINITY; - f64: 0.0, 1.0, f64::NEG_INFINITY, f64::INFINITY; - i64: 0, 1, i64::MIN, i64::MAX; -} - -/// The reduction `R` as a function from multisets of values to values, `Ω` in -/// the paper: the neutral element for the empty multiset, the element of a -/// multiset of one, and otherwise its elements combined in any order, as -/// `combine` is associative and commutative. -pub struct Reduce { - _marker: PhantomData (R, A)>, -} - -impl Reduce { - pub fn new() -> Self { - Self { - _marker: PhantomData, - } - } -} - -// Written out rather than derived, which would require `R: Default` and `A: Default`. -impl Default for Reduce { - fn default() -> Self { - Self::new() - } -} - -impl, A: Clone> Function for Reduce { - type Domain = Multisets>; - type Codomain = Values; - - fn apply(&self, multiset: &Multiset) -> A { - multiset.iter().cloned().fold(R::NEUTRAL, R::combine) - } - - fn domain(&self) -> Multisets> { - Multisets::new(Values::new()) - } - - fn codomain(&self) -> Values { - Values::new() - } -} - -pub type Sum = Reduce; -pub type Product = Reduce; -pub type Sup = Reduce; -pub type Inf = Reduce; - /// `Ω ∘ Ū ∘ ψ`, the reduction tensor of the paper, with reduction `R`, /// underlying tensor `U` and index function `M`: at each index, the reduction /// of the multiset of the values of `U` at the indices `M` maps it to. Built /// as `compose(Sum::new(), compose(multiset_lift(u), psi))`. -pub type ReductionTensor = Compose, M>>; +pub type ReductionTensor = Composition, M>>; diff --git a/tests/functions.rs b/tests/functions.rs index e3735d4..11918ae 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -1,13 +1,13 @@ use std::rc::Rc; use cten::{ - COO, Compose, Domain, Function, HRect, Physical, Stride, Sum, Tensor, ValueFunction, clamp, - compose, concatenate, diagonalize, eye, eye_rect, mm, movedim, multiset_lift, negate, ones, - permute, prod, scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, sum, - transpose, unsqueeze, where_, zeros, + COO, Composition, Function, HRect, IndexSet, Physical, Stride, Sum, Tensor, ValueFunction, + clamp, compose, concatenate, diagonalize, eye, eye_rect, mm, movedim, multiset_lift, negate, + ones, permute, prod, scale, sparse_coo_tensor, sparse_latticed_tensor, square, squeeze, stack, + sum, transpose, unsqueeze, where_, zeros, }; -type Strided = Compose>, Rc>; +type Strided = Composition>, Rc>; fn physical() -> Rc { Rc::new(Physical::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) @@ -202,7 +202,7 @@ fn transpose_is_an_involution() { assert_eq!(collect_2d(&restored, 3, 2), collect_2d(&*original, 3, 2)); } -// The mapping permutes indices rather than strides, so it applies to tensors +// The index function permutes indices rather than strides, so it applies to tensors // that have no strides to permute in the first place. #[test] fn transpose_applies_to_a_tensor_with_no_strides() { diff --git a/tests/mapping.rs b/tests/index_functions.rs similarity index 98% rename from tests/mapping.rs rename to tests/index_functions.rs index d7b5106..fa746a4 100644 --- a/tests/mapping.rs +++ b/tests/index_functions.rs @@ -1,9 +1,9 @@ use std::{assert_eq, collections::HashMap, rc::Rc}; use cten::{ - Affine, COO, Concat, Diagonal, Domain, EmptyMapping, HRect, InvAffine, InvProjection, - JaggedDomain, Mapping, Multiset, Permutation, Physical, Projection, Squeeze, Stride, Unsqueeze, - Where, + Affine, COO, Concat, Diagonal, EmptyMapping, HRect, IndexFunction, IndexSet, InvAffine, + InvProjection, JaggedDomain, Multiset, Permutation, Physical, Projection, Squeeze, Stride, + Unsqueeze, Where, }; /// The indices of a multiset in sorted order, to compare multisets whose diff --git a/tests/logic.rs b/tests/logic.rs index ed1c956..a3d7ab1 100644 --- a/tests/logic.rs +++ b/tests/logic.rs @@ -1,9 +1,9 @@ use std::rc::Rc; use cten::{ - Affine, COO, Diagonal, EmptyMapping, Function, HRect, IndexFunctionSignature, Inf, Mapping, - Multiset, Multisets, Physical, Product, Stride, Sum, Sup, Tensor, ValueFunction, compose, - multiset_lift, + Affine, COO, Diagonal, EmptyMapping, Function, HRect, IndexFunction, Inf, + MultiIndexFunctionSignature, Multiset, Multisets, Physical, Product, Stride, Sum, Sup, Tensor, + ValueFunction, compose, multiset_lift, }; fn physical() -> Rc { @@ -76,7 +76,7 @@ fn a_zero_stride_broadcasts_along_its_axis() { } } -// A mapping into flat positions needs a one-dimensional tensor under it. That +// An index function into flat positions needs a one-dimensional tensor under it. That // used to be a matter of index types; it is now checked when composing. #[test] #[should_panic(expected = "does not lie in")] @@ -106,7 +106,7 @@ fn left_composition_transforms_values_pointwise() { ); } -// A one-dimensional tensor is read by a mapping just like a physical one is: +// A one-dimensional tensor is read by an index function just like a physical one is: // there is no longer a wrapper between the two. #[test] fn a_sparse_vector_can_be_read_through_a_mapping() { @@ -341,17 +341,17 @@ fn an_affine_mapping_reads_a_lattice_out_of_a_tensor() { /// An index function sending each index to the multiset holding it twice. struct Twice { - signature: IndexFunctionSignature, + signature: MultiIndexFunctionSignature, } -impl Mapping for Twice { +impl IndexFunction for Twice { type Codomain = Multisets; fn map(&self, index: &[usize]) -> Multiset> { Multiset::new(vec![index.to_vec(), index.to_vec()]) } - fn signature(&self) -> &IndexFunctionSignature { + fn signature(&self) -> &MultiIndexFunctionSignature { &self.signature } } @@ -361,7 +361,7 @@ impl Mapping for Twice { #[test] fn a_reduction_counts_an_index_as_often_as_it_is_repeated() { let twice = Rc::new(Twice { - signature: IndexFunctionSignature::new( + signature: MultiIndexFunctionSignature::new( HRect::new(vec![6]), Multisets::new(HRect::new(vec![6])), ), diff --git a/tests/domain.rs b/tests/sets.rs similarity index 97% rename from tests/domain.rs rename to tests/sets.rs index 4b55026..b2fd0c0 100644 --- a/tests/domain.rs +++ b/tests/sets.rs @@ -1,4 +1,4 @@ -use cten::{Domain, HRect, JaggedDomain}; +use cten::{HRect, IndexSet, JaggedDomain}; #[test] fn a_hrect_holds_exactly_the_indices_within_its_shape() { diff --git a/tests/tensor.rs b/tests/tensor.rs index f9f5ede..10384a2 100644 --- a/tests/tensor.rs +++ b/tests/tensor.rs @@ -1,7 +1,7 @@ use std::rc::Rc; use cten::{ - COO, Compose, Diagonal, Domain, Function, HRect, InvProjection, JaggedTensor, Multiset, + COO, Composition, Diagonal, Function, HRect, IndexSet, InvProjection, JaggedTensor, Multiset, Multisets, Permutation, Physical, Stride, Sum, ValueFunction, Values, compose, multiset_lift, }; @@ -30,7 +30,7 @@ fn a_composition_takes_its_domain_from_the_function_applied_first() { assert_eq!(strided.domain().shape(), &[3, 2]); assert!(strided.domain().contains(&[2, 1])); assert!(!strided.domain().contains(&[2, 2])); - // The same storage under a mapping of another shape is another domain. + // The same storage under an index function of another shape is another domain. assert_eq!(transposed.domain().shape(), &[2, 3]); assert!(transposed.domain().contains(&[1, 2])); } @@ -74,7 +74,7 @@ fn a_function_on_values_keeps_the_domain_of_the_tensor_it_follows() { // Composition is not only for tensors: permuting the axes of indices one way and // then back is the identity on indices. #[test] -fn index_maps_compose_like_any_other_functions() { +fn index_functions_compose_like_any_other_functions() { let forth = Permutation::new(HRect::new(vec![4, 2, 3]), vec![1, 2, 0]); let back = Permutation::new(HRect::new(vec![2, 3, 4]), vec![2, 0, 1]); let round_trip = compose(back, forth); @@ -105,7 +105,7 @@ fn a_lifted_function_on_values_applies_to_each_element() { } /// The matrix of the paper's row sums example, `[[1, 2, 3], [0, 4, 1], [2, 2, 5]]`. -fn figure_matrix() -> Rc> { +fn figure_matrix() -> Rc> { Rc::new(compose( Physical::new(vec![1.0, 2.0, 3.0, 0.0, 4.0, 1.0, 2.0, 2.0, 5.0]), Stride::new(HRect::new(vec![3, 3]), vec![3, 1]),