MLX: Subtensor and IncSubtensor fail on a symbolic integer index under mx.compile
Nobody has claimed this yet.
Assessment
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Newbie friendliness
- 78/100
- Issue type
- Bug
- Clarity
- Clearly specified
- Activity status
- Active
- Tech stack
- python
- Domain
- backend, machine-learning
Research direction
Start with mlx_funcify_Subtensor and mlx_funcify_IncSubtensor, then run the provided symbolic-index reproducer under the MLX mode. Check both reads and writes with a traced scalar integer index; done means mx.compile accepts those operations without int() coercion while slice bounds retain their required handling.
Written by the indexing model from the issue text.
Description
x[i] and set_subtensor(x[i], y) with a symbolic integer i fail to compile on the MLX backend. mlx_funcify_Subtensor and mlx_funcify_IncSubtensor call int() on every integer index input, and under mx.compile the index is a traced mx.array, so int() raises. mlx accepts an integer mx.array as an index for both reads and writes, and the same graph compiles when the index is a one-element vector, because AdvancedSubtensor passes the array through.
import numpy as np
import pytensor
import pytensor.tensor as pt
x = pt.matrix("x")
i = pt.iscalar("i")
fn = pytensor.function([x, i], x[i], mode="MLX")
print(fn(np.eye(3, dtype="float32"), 1)) # ValueError: [eval] Attempting to eval an array during function transformations
# workaround: index with a one-element vector, x[i[None]][0], which lowers to AdvancedSubtensor
The int() coercion came in with #2240 for slice bounds. Restricting it to slice bounds and passing an integer mx.array index straight to x[index] would cover both ops.
- Dominant language
- Python
- Stars
- 644
- Forks
- 209
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from pymc-devs/pytensor
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
graph rewriting linalg performance
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 84/100
-
bug torch
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
All issues in pymc-devs/pytensor
Similar issues
-
Difficulty 2/5 1-3 hours Newbie friendliness 88/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
enhancement
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
-
Difficulty 2/5 1-3 hours Newbie friendliness 74/100