Generic matmul-based expansion for all-axis mean reduction

[AlekseiChirkovVention]

Context

While reviewing how downstream consumers adopt framework-owned dependency analysis, lowering, and optimizer tracing, we identified one piece of lowering logic that is not backend-specific and can be generalized in tinychain.autodiff instead of being reimplemented by every consumer.

The mean expansion rewrites an all-axis rank-2 MeanOperator into:

1. a MATMUL against a column of ones for row-sum reduction,
2. a MATMUL against a row of ones for column-sum reduction,
3. a scalar MUL by 1 / (rows * cols).

The core arithmetic uses matrix multiplication and multiplication. Opt-in expanded artifacts represent generated tensor-valued ones with a zero-input FillOperator and, for keepdims=False, represent the truthful [1, 1] to scalar rank transition with a trivial ReshapeOperator. These explicit operators prevent hidden operands and false output shapes.

Decision

The expanded-handler contract is accepted. Backends consuming opt-in expanded artifacts must support FillOperator in both tiers and trivial ReshapeOperator for the rank-reducing tier. MatmulOperator remains an ordinary two-operand operator. Existing non-expanded lowering contracts remain unchanged.

Proposal

Add reusable opt-in expansion passes to tinychain.autodiff, alongside lower_graph and lower_derivative_program, that expand a supported all-axis MeanOperator and the exact generated mean-VJP broadcast-scale region into matmul-based arithmetic with explicit constants and truthful rank changes.

Backends with a native reduce operation remain unaffected unless they explicitly run an expansion pass.

Acceptance criteria

• Public tinychain.autodiff APIs expand a supported all-axis mean and the exact generated mean-VJP broadcast-scale region for backends without a native reduce operation.
• Expansion is opt-in and off by default; existing non-expanded lowered output and handler requirements do not change.
• Expanded artifacts use a zero-input FillOperator for generated tensor-valued constants in both tiers.
• For keepdims=False, expanded artifacts use a trivial ReshapeOperator for the truthful [1, 1] to scalar rank transition.
• MatmulOperator remains an ordinary two-operand operator, and every emitted output_typespec matches the true operation result.
• A backend missing a handler required by an expanded artifact fails closed with unsupported_operator before any handler runs.
• Numerical equivalence between expanded and non-expanded lowering is validated against a generic limited-operation reference consumer used only for testing.

[AlekseiChirkovVention]

client-128-generic-matmul-based-mean-expansion.md

SDD spec for implementation of this issue

[AlekseiChirkovVention]

Decision A is accepted and recorded in the issue body. The current Decision and Acceptance criteria sections are authoritative and supersede the earlier draft criterion and the SDD attachment above.

Opt-in expanded artifacts require FillOperator support in both tiers and trivial ReshapeOperator support for keepdims=False. MatmulOperator remains an ordinary two-operand operator, emitted shapes remain truthful, and existing non-expanded lowering contracts remain unchanged.

[AlekseiChirkovVention]

client-128-generic-matmul-based-mean-expansion.md

Updated SDD document

[AlekseiChirkovVention]

Shipped in 229dfa5 — pull requests #130-#137, #139, #140.

tinychain.autodiff now exposes opt-in passes that rewrite a supported all-axis rank-2 mean
into two matmuls against generated constants plus a reciprocal-count multiply, and the
semantically matching broadcast-and-scale region of a derivative program into the same shape.
A backend restricted to {ADD, MUL, MATMUL} can lower a forward graph and the derivative
generated from it with no reduction, broadcast, or division handler.

Against the acceptance criteria

• Public passes expand a supported all-axis mean and the matching gradient region — done, as
  expand_mean_graph and expand_mean_derivative_program, plus *_detailed forms returning
  provenance records.
• Opt-in and off by default — opting in is calling a pass. No existing entry point gained a
  flag, keyword, or branch, and an artifact never passed through a pass lowers exactly as
  before.
• Zero-input FillOperator for generated constants in both tiers — done, as the single public
  representation, with a validated descriptor and one shared reader so no backend hand-parses
  operator parameters.
• Trivial ReshapeOperator for the keepdims=False rank transition — done, as a real reshape
  node rather than an elementwise node relabelled with a different rank.
• MatmulOperator stays an ordinary two-operand operator and every emitted output_typespec
  matches the true result — done, and the second half is stronger than it sounds: every emitted
  declared type is derived from its operands, never copied from the node it replaces. A
  copied type cannot be wrong, but it also cannot detect a wrong one.
• A backend missing a required handler fails closed with unsupported_operator before any
  handler runs — done, proved by a recording registry that stays empty, which itself asserts
  the registry holds other handlers so the empty record is not vacuous.
• Numerical equivalence against a generic limited-operation reference consumer — done, over
  both tiers, both differentiable dtypes, and square and non-square shapes: dtype and shape
  exactly equal, values within 1e-6 (f32) and 1e-12 (f64).

Worth knowing when adopting this

The rank-reducing tier needs a trivial reshape handler in both the forward artifact and the
derivative program. The gradient path's own leading seed reshape survives expansion untouched,
so registering one for the forward graph alone still fails closed on the derivative program.
Writing the loss with keepdims=True avoids the requirement entirely and is the recommended
lower-cost form.

The gradient-path rewrite is identified by what a region computes, not by where it came from:
no artifact carries an origin marker, so a claim to recognise "the generated mean gradient"
would be unprovable. Its exact-divisor clause is a scope guard, not a correctness guard — the
rewrite preserves value for any divisor, and the clause confines the pass to the reduction case
and keeps the one inexact substitution on an exact integer element count.

A pass makes no promise that its output stays differentiable. A consumer needing both a
derivative and an expanded artifact differentiates the unmodified source graph and expands only
the artifacts destined for analysis and lowering.

Not included, deliberately. The cross-specification composition test lands with #129, which
it depends on; the shared reference consumer and the expansion-hook signature it needs are
already on main. Partial-axis means, ranks other than 2, and expansion of other reductions
remain out of scope as designed.

Focused autodiff suite 910 passed; broad non-integration suite 1161 passed, with only the
pre-existing environment-dependent failures this repository already carried.