pymc-devs / pymc-devs/pytensor-ml
Guard the MLX convolution's rank at compile time
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 9
- Forks
- 7
- Avg merge
- 6h 55m
- Merged PRs (30d)
- 40
Description
dispatch/mlx/pooling.py guards the rank it can handle and raises at compile time with a message naming the backend. dispatch/mlx/conv.py has no such guard, so a rank-4 convolution compiles cleanly and fails inside MLX on the first call, with a message that names neither the layer nor the backend policy. Every other NotImplementedError in the dispatch tree fires during pytensor.function(); this is the only one that waits until data loading, initialization and mx.compile are all behind you.
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor_ml.layers.conv import ConvLayer
X = pt.tensor("X", shape=(1, 4, 4, 4, 4, 2))
W = pt.tensor("W", shape=(2,) * 6)
f = pytensor.function([X, W], ConvLayer((2,) * 4, (1,) * 4, (1,) * 4)(X, W), mode="MLX")
# compiles fine
f(np.zeros((1, 4, 4, 4, 4, 2)), np.zeros((2,) * 6))
# ValueError: [conv] Only works for inputs with 1-3 spatial dimensions.
One rank check mirroring the one already in the sibling pooling module turns this into a compile-time error like every other backend hole.
Contributor guide
No contributing guide indexed for this repository
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.
Research direction
Compare the rank guard in dispatch/mlx/pooling.py with dispatch/mlx/conv.py, then use the issue's pytensor.function() reproduction as the first check. Done means unsupported convolution ranks fail during compilation with a message naming the layer and backend policy, rather than on the first call.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 88/100