Skip to content

Fix sparse.dot gradient for 1d dense operands#2287

Open
jarotter wants to merge 1 commit into
pymc-devs:mainfrom
jarotter:sparse_matrix_to_vector_dot
Open

Fix sparse.dot gradient for 1d dense operands#2287
jarotter wants to merge 1 commit into
pymc-devs:mainfrom
jarotter:sparse_matrix_to_vector_dot

Conversation

@jarotter

Copy link
Copy Markdown

Description

pytensor.sparse.dot supports a 1d dense operand — Dot.make_node accepts it, infer_shape handles it, and the forward pass computes correctly (covered by the existing test_csr_dense/test_csc_dense tests). But Dot.pullback assumed both operands were always 2d, so differentiating anything containing sparse.dot(X, beta) (or sparse.dot(beta, X)) with a 1d beta raised inside pytensor.grad itself, 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:

ValueError: Incompatible shared dimension for dot product: (1, 20), (5, 1)

or

TypeError: The dense dot product is only supported for dense types

Both are the same underlying defect, just diverging on which branch of pullback hits 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 promotion dense_dot already does for the dense Dot op. Behavior for the already-working 2d/2d case is untouched (verified with dprint tcase is byte-identical to before).

Testing

Added 8 tests to TestDots in tests/sparse/test_math.py:

  • test_dot_vector_operand_grad — the core fix, both operand ord
  • test_dot_vector_operand_grad_length_one — broadcastable lengt
  • test_dot_vector_operand_grad_shape_regimes — both failure-modlved shape)
  • test_dot_vector_operand_grad_squared_error_loss — end-to-end rted pattern
  • test_dot_matrix_sparse_operand_grad — non-regression: grad w.e already-working 2d case, previously untested

Related Issue

Checklist

  • Checked that the pre-commit linting/style checks pass
  • Included tests that prove the fix is effective or that the new feature works
  • Added necessary documentation (docstrings and/or example notebooks)
  • If you are a pro: each commit corresponds to a [relevant logical change]

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BUG: grad of sparse.dot fails when the operand is a 1d vector

2 participants