♻️ refactor(transforms): declare groups() on AbstractTransform - #730
♻️ refactor(transforms): declare groups() on AbstractTransform#730nstarman wants to merge 3 commits into
groups() on AbstractTransform#730Conversation
`_merge` had rules for `Rotate` and `AbstractAdd` only, so `simplify(S1 | S2)` returned the two-element `Composed` unchanged and a chain of scalings paid one full `act` dispatch per element. A three-scaling chain cost 8822 us eagerly against 3953 us fused, 2.2x. The merge is a matrix product, so it is exact -- the fused and sequential results agree to 0.0, not merely to a tolerance -- and it runs regardless of `approx`, like the `Rotate` rule it mirrors. Mixed linear chains (`Rotate | Scale`) still do not fuse; that is 4.9x and needs a return type this library does not yet have. Filed separately. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`simplify` could merge two rotations or two scalings, but not a rotation with a scaling: the product is neither, and there was no type to put it in. A five-operator linear chain therefore stayed five operators, paying a full `act` dispatch each -- 11481 us eagerly against 3415 us fused, 3.4x, with the two agreeing to 8.9e-16. `Linear` is a sibling `@final` under `AbstractLinearTransform`, which already owns the matrix, its validation and every `act` path, so the class is a field and three members. The generic `_merge` rule is the fallback; the same-type rules are more specific and still win, so `Rotate | Rotate` keeps returning a `Rotate`. Nothing that fused before fuses differently. The one new idea is that the group travels with the matrix rather than with the type: `Rotate | Reflect` is still orthogonal, and reporting it as merely affine would lose that. `Composed.groups` already computes exactly this least common supergroup, and group classes are registered static, so it is a field. Mismatched dimensions -- a 3x3 `Rotate` beside a 4x4 `LorentzBoost` -- decline to merge. That guard is tested through `_merge` rather than `simplify`, which has no rule registered for `LorentzBoost` and raises first; pre-existing, fixed separately. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Every concrete transform defines `groups()`, and callers depend on it -- `Composed.groups` reads it off each component, `_merge` off both operands -- but the base declared nothing, so it was a protocol by convention. A type checker could only see the attribute where the caller happened to hold a concrete type: `ty` passes `Composed.groups` because its `transforms` field is loosely typed, and fails the moment a parameter is annotated `AbstractLinearTransform`. That is what the `ty: ignore` in `general_linear` was covering; it is removed here. Declared as an instance method so both existing forms satisfy it. Most operators answer from their type alone and implement it as a `classmethod`, which stays reachable from an instance; `Composed` and `Linear` compute it from what they hold. Abstract rather than a default: all eleven concrete transforms already implement it, and only the abstract intermediates (`AbstractAdd`, `AbstractLinearTransform`, `AbstractCompositeTransform`) do not, which is what abstract means. A default returning some fallback group would be a wrong answer waiting for a subclass that forgot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR strengthens the transforms API by formally declaring groups() on AbstractTransform, making group membership an explicit part of the base interface while supporting both class-derived and instance-computed implementations. It also expands simplification/fusion capabilities for linear transforms by introducing a general Linear operator and adding additional _merge rules with accompanying tests.
Changes:
- Add an abstract
groups()method toAbstractTransformto formalize group membership as part of the core transform interface. - Introduce
Linear(a general linear transform that carries its group as data) and a generic_merge(AbstractLinearTransform, AbstractLinearTransform)fallback for mixed linear fusion. - Add a
_merge(Scale, Scale)rule and extendsimplifytests to cover new scale merging and mixed linear fusion behavior.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/transforms/test_simplify.py | Adds tests for Scale merging and mixed linear fusion into Linear, including group propagation assertions. |
| src/coordinax/transforms/_src/actions/scale.py | Adds _merge(Scale, Scale) to fuse adjacent scales via matrix multiplication. |
| src/coordinax/transforms/_src/actions/general_linear.py | Adds the new Linear transform, its group-carrying behavior, simplify rule, and generic linear _merge fallback. |
| src/coordinax/transforms/_src/actions/base.py | Declares AbstractTransform.groups() as an abstract method with documentation and example. |
| src/coordinax/transforms/_src/actions/init.py | Exports the new general_linear action module. |
| src/coordinax/transforms/init.py | Re-exports Linear as part of the public coordinax.transforms API. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| def __init__( | ||
| self, M: Any, group: type[groups.AbstractTransformGroup] = groups.AffineGroup | ||
| ) -> None: | ||
| object.__setattr__(self, "M", jnp.asarray(M)) | ||
| object.__setattr__(self, "group", group) |
| "Translate", | ||
| "LorentzBoost", | ||
| "Linear", | ||
| "Rotate", | ||
| "Reflect", |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #730 +/- ##
=======================================
Coverage 96.54% 96.55%
=======================================
Files 265 266 +1
Lines 8780 8828 +48
=======================================
+ Hits 8477 8524 +47
- Misses 303 304 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Stacks on #728 — review the top commit only, or merge #728 first. It removes the
ty: ignorethat PR had to add.The gap
Every concrete transform defines
groups(), and callers depend on it —Composed.groupsreads it off each component,_mergeoff both operands — butAbstractTransformdeclared nothing. It was a protocol by convention, with nothing on the base to point at.A type checker could therefore only see the attribute where the caller happened to hold a concrete type.
typassesComposed.groupsmerely because itstransformsfield is loosely typed, and fails the moment a parameter is annotatedAbstractLinearTransform— which is exactly what #728 hit, and what itsty: ignorewas covering.The fix
One abstract method on the base. Spelled as an instance method, so both existing forms satisfy it:
classmethod, which stays reachable from an instanceComposedandLinearcompute it from what they holdAbstract rather than a default. All eleven concrete transforms already implement it:
The three that don't are the abstract intermediates, which is what abstract means. A default returning some fallback group would be a wrong answer waiting for a subclass that forgot to override — and since
groups()feedsleast_common_supergroup, a wrong group silently widens every fusion that touches it.Verification
prek run --all-filesclean, includingtywith ✨ feat(transforms): addLinear, so mixed linear chains fuse #728's suppression removed🤖 Generated with Claude Code