Fix sparse.dot gradient for 1d dense operands#2287
Open
jarotter wants to merge 1 commit into
Open
Conversation
Dot.pullback assumed both operands of a sparse dot were always 2d. This broke whenever the dense operand was a plain vector - e.g. sparse.dot(X, beta) with a 1d coefficient vector beta, the standard way to multiply a sparse design matrix against a dense coefficient vector. The forward pass already supported this shape; only the gradient was broken, so every gradient-based sampler failed as soon as pytensor.grad was called. The fix is to promote the 1d operand, and the incoming gradient to a 2d row/column before building the gradient, reuse the existing matrix-product formulas unchanged, and squeeze the added axis back out of the result. Behavior for the already-working 2d/2d case is unchanged. Fixes pymc-devs#2282
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
pytensor.sparse.dotsupports a 1d dense operand —Dot.make_nodeaccepts it,infer_shapehandles it, and the forward pass computes correctly (covered by the existingtest_csr_dense/test_csc_densetests). ButDot.pullbackassumed both operands were always 2d, so differentiating anything containingsparse.dot(X, beta)(orsparse.dot(beta, X)) with a 1dbetaraised insidepytensor.graditself, before any compilation.Depending on whether the sparse operand's static shape was known at graph-construction time, this surfaced as one of two different errors:
or
Both are the same underlying defect, just diverging on which branch of
pullbackhits a shape check first.Fix: promote the 1d operand, and the incoming gradient
gz, to a 2d row/column before building the gradient, reuse the existing matrix-product gradient formulas unchanged, then squeeze the added axis back out of the result — mirroring the promotiondense_dotalready does for the denseDotop. Behavior for the already-working 2d/2d case is untouched (verified withdprinttcase is byte-identical to before).Testing
Added 8 tests to
TestDotsintests/sparse/test_math.py:test_dot_vector_operand_grad— the core fix, both operand ordtest_dot_vector_operand_grad_length_one— broadcastable lengttest_dot_vector_operand_grad_shape_regimes— both failure-modlved shape)test_dot_vector_operand_grad_squared_error_loss— end-to-end rted patterntest_dot_matrix_sparse_operand_grad— non-regression: grad w.e already-working 2d case, previously untestedRelated Issue
gradofsparse.dotfails when the operand is a 1d vector #2282Checklist
Type of change