pymc-devs / pymc-devs/pytensor-ml

Guard the MLX convolution's rank at compile time

Open Beginner friendly
#135 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement pytensor
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

  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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.