Skip to content

Add rewrite with Theorem 1 - #20

Open
ppraneth wants to merge 1 commit into
mainfrom
theorem
Open

ppraneth wants to merge 1 commit into
mainfrom
theorem

Conversation

@ppraneth

@ppraneth ppraneth commented Nov 9, 2026

Copy link
Copy Markdown
Collaborator
  • Adds rewrite(tensor), which turns a tensor into an equivalent one that is cheaper to evaluate, or gives it back unchanged as Err when no identity applies.
  • For now it uses Theorem 1, f ∘ Σ ∘ Ū ∘ ψ = Σ ∘ (f ∘ U)‾ ∘ ψ, so a zero-preserving function like square or relu is only applied to the stored values. It checks that f(0) = 0 and that the tensor is sparse.
  • Adds a Sparsity trait so that index functions (COO, Diagonal, EmptyMapping, InvProjection, InvAffine) can tell whether they send each index to at most one index.

@ppraneth
ppraneth requested a review from ValerianRey November 9, 2026 04:04
@ValerianRey

Copy link
Copy Markdown
Member

I'll read in details later but I just have a quick comment so far. I don't think we should return an error when the rewrite is impossible. The idea is that rewrite will be identity when nothing can be improved, so that users (or even us) can call rewrite without a priori knowledge of the rewritability of the tensor, and let cten check that for them.

@ppraneth

ppraneth commented Nov 9, 2026

Copy link
Copy Markdown
Collaborator Author

Cool, I checked what Rust allows:

  1. When the theorem's checks fail (f(0) != 0 or the tensor isn't sparse), rewrite can give the tensor back instead of an error. I tried it: rewrite returns a small type that is either the original or the rewritten tensor, and it evaluates, prints and composes like any tensor. One catch: you can't rewrite that result again (e.g. scale it and rewrite), because its type isn't a plain sum tensor anymore. With the current Result you can.
  2. Calling rewrite on any tensor at all (a Physical, a row sum...) is harder. Right now rewrite only exists for tensors shaped like f ∘ Σ ∘ Ū ∘ ψ, because the type of the result depends on the shape. Adding "identity for every other shape" next to it doesn't compile (conflicting implementations), and the feature that would allow it (specialization) isn't in stable Rust. The way around it I see is to have rewrite return an Rc<dyn Tensor> and look at the composition at runtime, but that's a bigger change and I'd have to prototype it first.

Should I proceed as a whole or just do 1?

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.

2 participants