Skip to content

Add multiset-valued tensors - #19

Merged
ppraneth merged 4 commits into
mainfrom
multiset
Nov 6, 2026
Merged

ppraneth merged 4 commits into
mainfrom
multiset

Conversation

@ppraneth

@ppraneth ppraneth commented Nov 5, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #18

  • Adds Multiset, stored as an array and printed as {{1, 3, 3}}.
  • Everything is now a Function with a real domain and codomain (index sets, Values<A> or Multisets<S>), and compose and multiset_lift replace RightComposition, LeftComposition, BoxedTensor and the ReductionTensor struct.
  • Reductions reduce multisets and index functions return them, so a sum tensor is compose(Sum::new(), compose(multiset_lift(u), psi)) and repeated indices count.

@ppraneth
ppraneth requested a review from ValerianRey November 5, 2026 08:01

@ValerianRey ValerianRey left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we get rid of the Bar class and instead have just a function (e.g. multiset_lift) that will build a multiset-index multiset-valued tensor? I guess this requires either making tensors generic on their indexing too (be able to take multiset of indices instead of just normal indices), or making two kinds of tensors, those indexed by normal indices and those indexed by multisets of indices.

Can we get rid of BarComposition entirely, and instead just use LeftComposition of a reduction and a multiset-indexed multiset-valued tensor?

Can we get rid of ReductionTensor, that I think is now just a wrapper around a BarComposition?

I'm not sure about exactly what is possible among that, but the core idea is to avoid having too many classes.

@ValerianRey

Copy link
Copy Markdown
Member

In a sense the very concept of Tensor could maybe be removed too. The building brick is really functions, from a domain (possibly multisets) to a co-domain (possibly multisets), compositions between those functions (also, we don't need left composition and right composition, we just need a compose function that takes two functions, and we can change from left composition to right composition by swapping the arguments when calling this compose).
And lastly we need the multiset_lift function to transform a function from A to B to a function from the set of multisets of A to the set of multisets of B.

Not sure if these should be functions, classes or whatnot (not good enough at rust to know that), but that's the idea.

So I guess this PR actually requires very deep changes to the whole library (almost an entire rewrite). But in the end it can be 5x simpler and more powerful.

What do you think?

@ppraneth

ppraneth commented Nov 5, 2026 •

Copy link
Copy Markdown
Collaborator Author

Reworked it. Now there's one Function trait (input, output, apply) and everything implements it: tensors, index maps, index functions, reductions. compose(f, g) replaces RightComposition, LeftComposition and BarComposition, and multiset_lift replaces Bar. ReductionTensor isn't a class anymore, a sum tensor is just compose(Sum::new(), compose(multiset_lift(u), psi)) like in the paper. The 6 mapping traits became one too.

I kept Tensor as a trait alias (any function from indices to values) so we can still write Rc<dyn Tensor>, and closures need a small ValueFunction wrapper because of how Rust traits work. Functions on indices still carry their domain and compose checks they match, functions on values just don't have one.

It's not 5x less code since most of it is the mappings themselves, but there are way fewer concepts now. What do you think?

@ppraneth

ppraneth commented Nov 5, 2026

Copy link
Copy Markdown
Collaborator Author

A few things I kept on purpose, and why:

  • Tensor and BoxedTensor are still there but only as aliases, nothing implements them directly. A tensor is just any function from indices to values. I kept them so we can write Rc<dyn Tensor> instead of Rc<dyn Function<Input = [usize], Output = f32>>.
  • ReductionTensor is only a type alias now, to keep return types short.
  • compose and multiset_lift still return small structs (Compose, MultisetLift) because in Rust a function has to return some type. I also had to add ValueFunction to wrap closures and Reduce for the reductions.
  • Boxed stays so only tensors on a hrect can be printed, and that's checked at compile time.
  • Functions on values have no domain, and functions on multisets of indices report the set the indices come from.

About the 5x: most of the code is the mappings, the torch functions, printing and the lattice solver, and none of that changed, so the core is about the same size, just with fewer concepts.

Two things we could do to go further:

  1. Remove the Tensor, BoxedTensor and ReductionTensor aliases. Small cut, but types get longer to write.
  2. Turn the 8 mappings that are really affine (Stride, Permutation, Projection, Unsqueeze, Squeeze, InvProjection, Diagonal, EmptyMapping) into small constructors over Affine / InvAffine, leaving 5 kinds of index functions (Affine, InvAffine, COO, Concat, Where). That's the big one, about 230 lines less. The downside is that the diagonal and axis sums go through the general lattice solver instead of a direct formula, and the error messages get more generic.

Should I do one or both, and in this PR or a separate one? What do you think?

@ppraneth
ppraneth requested a review from ValerianRey November 5, 2026 10:20
@ValerianRey

Copy link
Copy Markdown
Member
  • Tensor and BoxedTensor are still there but only as aliases, nothing implements them directly. A tensor is just any function from indices to values. I kept them so we can write Rc<dyn Tensor> instead of Rc<dyn Function<Input = [usize], Output = f32>>.
  • ReductionTensor is only a type alias now, to keep return types short.
  • Boxed stays so only tensors on a hrect can be printed, and that's checked at compile time.

Perfect.

  • compose and multiset_lift still return small structs (Compose, MultisetLift) because in Rust a function has to return some type. I also had to add ValueFunction to wrap closures and Reduce for the reductions.

I think the Compose class is necessary, but idk about the MultisetLift class. Isn't it just a Function whose domain is a multiset and whose codomain is a multiset too? Any way to replace the return type of multiset_lift by that?

  • Functions on values have no domain, and functions on multisets of indices report the set the indices come from.

What are functions on values? Why don't they have a domain? I don't understand. Also I'm not sure what you mean by "functions on multisets of indices report the set the indices come from".

  1. => Let's not do that
  2. => We can do that in a future PR. I feel like some of these mappings should inherit (or the rust equivalent) from Affine, so that we can specialize their resolution without going through the lattice solver. For now let's just keep them separate.

@ppraneth

ppraneth commented Nov 5, 2026

Copy link
Copy Markdown
Collaborator Author

Isn't it just a Function whose domain is a multiset and whose codomain is a multiset too? Any way to replace the return type of multiset_lift by that?

Yeah Rust can do that. multiset_lift can return impl Function<Input = Multiset<...>, Output = Multiset<...>>, so callers only see a function from multisets to multisets and MultisetLift becomes private. The catch is that the ReductionTensor alias names MultisetLift, so it would have to go, and the functions that build reduction tensors (sum, prod, mm, diagonalize, zeros, ones, eye...) would return impl BoxedTensor<A> + Display instead. That's actually shorter than the types we have now. Should I do it?

What are functions on values? Why don't they have a domain?

Functions on values are things like square, x > 3, or the reductions, they take values instead of indices. They do have a domain, it's just the Rust type (f32 for example), and the compiler already checks it when composing (composing square after an index map won't compile). domain() only returns index sets because a tensor's shape isn't part of its type, so that one has to be checked when running.

I'm not sure what you mean by "functions on multisets of indices report the set the indices come from".

The domain of multiset_lift(u) is really the multisets of u's indices. In the code it just returns u's index set, because that's all the check needs: when composing with an index function, it compares that set with the indices the index function returns.

1 and 2 sound good, keeping the aliases and leaving the affine stuff for a future PR.

@ValerianRey

Copy link
Copy Markdown
Member

Yeah Rust can do that. multiset_lift can return impl Function<Input = Multiset<...>, Output = Multiset<...>>, so callers only see a function from multisets to multisets and MultisetLift becomes private. The catch is that the ReductionTensor alias names MultisetLift, so it would have to go, and the functions that build reduction tensors (sum, prod, mm, diagonalize, zeros, ones, eye...) would return impl BoxedTensor<A> + Display instead. That's actually shorter than the types we have now. Should I do it?

Can't we just have also a Function<Input = Multiset<...>, Output = Multiset<...>> in the ReductionTensor alias?
Then, the question is about whether we want an alias to avoid having such a long type name repeated all over the code, and I think it would be a good thing. But I don't like the name MultisetLift, I would simply call it MultiFunction maybe? As in, a MultiFunction is a multiset-domain multiset-codomain function.

domain() only returns index sets

Seems like a core issue.

The domain of multiset_lift(u) is really the multisets of u's indices. In the code it just returns u's index set, because that's all the check needs: when composing with an index function, it compares that set with the indices the index function returns.

Seems very hacky too.

Is there any way to fix those last two points?

@ppraneth

ppraneth commented Nov 5, 2026

Copy link
Copy Markdown
Collaborator Author

Can't we just have also a Function<Input = Multiset<...>, Output = Multiset<...>> in the ReductionTensor alias?

Not really in Rust. A type alias can only name a concrete type, or a dyn trait object which adds boxing and slower calls. Writing "some Function from multisets to multisets" in an alias isn't in stable Rust yet. So I'd keep the struct, rename it to MultiFunction like you said, and keep it in the alias.

Is there any way to fix those last two points?

Yes, I tried it in a local copy and it works. The idea is that every domain and codomain becomes a real set:

  • indices: HRect / JaggedDomain like now
  • values: a new Values<A>, e.g. Values<f32> is ℝ
  • multisets: a new Multisets<S>, which is M(S)

So domain() and codomain() always return a set, no more None. square goes from Values<f32> to Values<f32>, a reduction from Multisets<Values<f32>> to Values<f32>, multiset_lift(u) from Multisets<HRect> (the multisets of u's indices) to Multisets<Values<f32>>, and an index function from HRect to Multisets<HRect>, which is literally ψ: V → M(U). compose then checks that the codomain of the inner function lies in the domain of the outer one, for every kind of set.

Some side effects:

  • Boxed and BoxedTensor go away, because "boxed" just means the domain is an HRect now. Tensor<A> has HRect as its default domain, so printing is still checked at compile time.
  • The compose check is a bit stricter than before, it checks that the shape fits and not only the number of coordinates. No existing test broke, one panic message changed.
  • .hrect() becomes .domain().

All tests pass (plus 2 new ones for the domains) and the demo output is the same.

Is it okay if I go ahead with this?

@ValerianRey

Copy link
Copy Markdown
Member

Sounds perfect, please go ahead!

@ppraneth

ppraneth commented Nov 5, 2026

Copy link
Copy Markdown
Collaborator Author

Done, pushed it. Every function has a real domain and codomain now, and MultisetLift is renamed to MultiFunction.

@ValerianRey ValerianRey left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very good, this is exactly what I had in mind. I think a few names can be a bit confusing though.

Could we maybe move Domain classes to sets.rs? IIUC Domains are a type of Sets, right?

Also, I think we should rename Domain to IndexSet.

Mapping (in mapping.rs) => IndexFunction.
IndexFunctionSignature => MultiIndexFunctionSignature
(and file rename, and other linked renames)

After that you can merge!

Comment thread src/function.rs Outdated
}

/// `f ∘ g`: applies `g`, then `f` to the result.
pub struct Compose<F, G> {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rename to Composition? it would make sense that the function is called compose but the class is called Composition.

Comment thread src/tensor.rs Resolved
@ppraneth
ppraneth merged commit 8c11e3d into main Nov 6, 2026
4 checks passed
@ppraneth
ppraneth deleted the multiset branch November 6, 2026 03:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add multiset-valued tensors

2 participants