f-dangel / curvlinops / 28752829889

92%

Build Type

push

github

Committed by web-flow

Commit Message

[ADD] Detect Linear weights/biases reused across reshape-wrapped matmuls (#308)

* [ADD] Detect Linear weights/biases reused across reshape-wrapped matmuls

Make the FX layer-IO matcher robust to graphs that wrap a Linear's matmul in
transparent, last-dim-preserving shape ops and add the bias only after them --
e.g. Taylor-mode/jet graphs that emit `mm -> _unsafe_view -> squeeze -> add(bias)`
and reuse one weight across several such matmuls (one per Taylor coefficient).

- linear.py: traverse last-dim-preserving transparent ops (view/_unsafe_view/
  reshape/squeeze/unsqueeze) when linking a bias `add` to its matmul, on both the
  weight and bias sides, keeping the two infos identical so they deduplicate.
  Reshapes that alter the last (feature) dimension are still rejected.
- verification.py: treat a usage path as covered when a detected path is a prefix
  of it (the usage reached a matched layer output and merely continues downstream),
  not only when the usage is a prefix of a detected path. Without this, a weight
  reused in matmuls of uneven reshape depth was spuriously reported as uncovered.

Existing behavior is unchanged for standard nn.Module graphs (io_collector and
computer test suites pass). No new public API.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [TEST] Cover reshape-wrapped bias detection in the FX layer-IO matcher

Add tests for the pattern the previous commit enables -- a Linear whose bias is
added after transparent, last-dim-preserving reshape ops, and a single weight
reused across matmuls of uneven reshape depth (as emitted by collapsed
Taylor-mode / jet graphs, one matmul per Taylor coefficient).

- test_param_io.py:
  - test_bias_add_through_transparent_reshapes: bias reached through
    `mm -> unsqueeze -> squeeze -> add(bias)` is linked to its layer.
  - test_weight_reused_across_reshape_wrapped_matmuls: one weight over three
    coefficient matmuls (only the forward one biased, via a reshape cha... (continued)

Coverage Stats

34 of 36 new or added lines in 1 file covered. (94.44%)

1 existing line in 1 file now uncovered.

2667 of 2888 relevant lines covered (92.35%)

0.92 hits per line