pymc-devs / pymc-devs/pytensor

pt.flatten drops static shape on a fully-known input

Open
#2,366 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug shape inference
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

flatten builds its reshape target as (*x.shape[:ndim-1], -1), and unpacking a slice of the shape vector gives entries that don't constant-fold, so the output type loses shapes an equivalent reshape keeps.

import pytensor.tensor as pt

x = pt.tensor("x", shape=(32, 24, 24, 16))

print(pt.flatten(x, ndim=2).type.shape)        # (None, None)
print(x.reshape((x.shape[0], -1)).type.shape)  # (32, 9216)

Potential fix (requires testing):

dims = (*(_x.shape[i] for i in range(ndim - 1)), -1)

Indexing per axis yields Shape_i, which folds where the extent is known; a slice of the shape vector yields a Subtensor that doesn't. This recovers (32, 9216) but not the partly-dynamic case — with an unknown batch, Reshape still can't resolve the -1 even when every other extent is known, so it stays (None, None) where (None, 9216) is derivable.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the pt.flatten entry point and compare its reshape-target construction with the equivalent x.reshape call shown in the issue. Verify the fully-known case produces (32, 9216), and check the partly-dynamic case to determine whether the known trailing extent can be preserved without changing the unresolved batch dimension.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.