[avantikalal]

Summary

• gReLU models pooled to a single value per task return shape (N, tasks, 1), not (N, tasks). tangermeme.deep_lift_shap selects a single task via y[:, target], which only consumes the tasks axis and leaves the trailing 1 in place.
• That shape mismatch between its output-difference tensor (N/2, 1) and its attribution-sum tensor (N/2,) is then silently broadcast by PyTorch into a meaningless (N/2, N/2) matrix instead of raising a shape error — corrupting deep_lift_shap's own internal convergence-delta sanity check and producing spurious/inflated "Convergence deltas too high" warnings.
• Add _SqueezeOutput, a small nn.Module adapter that drops a trailing length-1 axis before the model output reaches deep_lift_shap. Models with genuine multi-position outputs ((N, tasks, L>1)) are left untouched.
• This is not Enformer-specific — it affects any ConvHead-based model using the default pooling, so it's independent of Enable deepshap for enformer models #56/Enable DeepSHAP for Enformer models #201.

Fixes #202

Verification

• Verified against a real trained Enformer-based checkpoint (Genentech/human-atac-catlas-model): the same convergence deltas that previously appeared as a bogus (N, N)-shaped matrix with values up to ~10 now come out correctly shaped (matching the expected per-reference batch sizes) and roughly an order of magnitude smaller.
• Added test_squeeze_output, a deterministic unit test of the wrapper (with and without a trailing singleton dim), since the existing test_get_attributions model's warning behavior turned out to be data-dependent and didn't reliably discriminate the bug on its own.
• Extended test_get_attributions to assert no "Convergence deltas too high" warning fires.
• Full tests/test_interpret.py suite passes (12 passed, no regressions).

Test plan

• pytest tests/test_interpret.py passes
• test_squeeze_output passes and fails if the wrapper is removed (verified by temporarily reverting the score.py change)
• Confirmed on a real Enformer checkpoint that convergence deltas are correctly shaped after the fix

🤖 Generated with Claude Code