Skip to content

Fix GetItemList setting an incorrect static output shape#2275

Open
zlee-cf wants to merge 2 commits into
pymc-devs:mainfrom
zlee-cf:fix-2274-getitemlist-static-shape
Open

Fix GetItemList setting an incorrect static output shape#2275
zlee-cf wants to merge 2 commits into
pymc-devs:mainfrom
zlee-cf:fix-2274-getitemlist-static-shape

Conversation

@zlee-cf

@zlee-cf zlee-cf commented Jul 7, 2026

Copy link
Copy Markdown

Description

sparse.get_item_list(x, index) selects rows of a sparse matrix, so its output has len(index) rows. But GetItemList.make_node returned the output as x.type(), which clones x's full static shape (including its row count).

For a SparseConstant, that shape is fully specified (e.g. (500, 40)), so the output was statically typed with the parent's row count regardless of the index length. A symbolic (shared) matrix has shape (None, None), so it never hit this. infer_shape and the runtime value were always correct — only the static output type was wrong.

This broke vectorize_graph which trusts the static shape, inserts a SpecifyShape, and then asserts at runtime:

AssertionError: SpecifyShape: dim 0 of input has shape 5, expected 500.

The fix builds the output type with the row dim taken from the index (None when unknown, or its statically-known length) and the column dim preserved from x, matching infer_shape.

Related Issue

Checklist

Type of change

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

GetItemList.make_node returned x.type(), cloning x's full static shape. For
a SparseConstant the row count is fully specified, so the output was typed
with the parent's rows regardless of the index length, breaking
vectorize_graph via a spurious SpecifyShape. Take the row dim from the index
and keep x's column dim, matching infer_shape.
@ricardoV94 ricardoV94 force-pushed the fix-2274-getitemlist-static-shape branch from 1a58ad8 to 8805b8c Compare July 7, 2026 12:33
@ricardoV94

Copy link
Copy Markdown
Member

I fixed a few more wrong make_nodes

@ricardoV94 ricardoV94 added bug Something isn't working shape inference labels Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working shape inference

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BUG: sparse.get_item_list on a SparseConstant sets an incorrect static output shape, breaking vectorize_graph

2 participants