pymc-devs / pymc-devs/pytensor

Disable runtime broadcasting in indexing operations

Open
#1,348 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

Description

We are inconsistent in what Ops we allow runtime broadcasting and which we don't.

import pytensor.tensor as pt

x = pt.vector("x", shape=(None,))  # Not known to have length 1 at runtime
out = pt.alloc(x, 3, 5)
try:
    out.eval({x: [1]})
except Exception as e:
    print(str(e).splitlines()[0])
# Runtime broadcasting not allowed. The output of Alloc requires broadcasting a dimension of the input value, which was not marked as broadcastable. If broadcasting was intended, use `specify_broadcastable` on the relevant input.

out = pt.zeros((10, 10))[[5, 6, 7], [0, 1, 2]].inc(x)
try:
    out.eval({x: [1]})
except Exception as e:
    print(str(e).splitlines()[0])
else:
    print("Did not raise")
# Did not raise

Note that whenever we allow runtime broadcasting will have a wrong gradient wrt to that broadcasted input, since we never implemented a mechanism to reduce runtime broadcasted dimensions.

print(pt.grad(out.sum(), x).eval({x: [1]}).shape)  # (3,)

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 by reproducing the two indexing examples in the issue and tracing the indexing operation implementations and their gradient handling. Compare their runtime broadcasting checks with pt.alloc, then verify that indexing operations reject unmarked runtime broadcasting and no longer produce an incorrect gradient for that case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.