pytensor-ml 0.1.0__tar.gz → 0.2.0__tar.gz

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/.github/workflows/mypy.yml

@@ -25,8 +25,8 @@ jobs:
         shell: bash -l {0}
 
     steps:
-    - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
-    - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
       env:
         CACHE_NUMBER: 0
       with:

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/.github/workflows/pre-commit.yml

@@ -11,6 +11,6 @@ jobs:
     env:
       SKIP: no-commit-to-branch
     steps:
-    - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
-    - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
     - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/.github/workflows/release.yml

@@ -18,12 +18,12 @@ jobs:
       attestations: write
       id-token: write
     steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
           # hatch-vcs derives the version from git tags, so the full history is required.
           fetch-depth: 0
           persist-credentials: false
-      - uses: hynek/build-and-inspect-python-package@d44ca7d91762de7a7d5436ddae667c6da6d1c3df # v2.18.0
+      - uses: hynek/build-and-inspect-python-package@2abe76da66d0a6a4a227101f9348ee855797cfa5 # v3.0.1
         with:
           # Prove that the packages were built in the context of this workflow.
           attest-build-provenance-github: true
@@ -47,4 +47,4 @@ jobs:
           name: Packages
           path: dist
       - name: Publish Package to PyPI
-        uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0
+        uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/.github/workflows/run_tests.yml

@@ -20,17 +20,26 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         python-version: ["3.12"]
+        # Off by default, flipped to 1 by the per-backend jobs under `include`. A core job installs no
+        # backend at all, which is what proves the library runs without one.
+        install-jax: [0]
+        install-mlx: [0]
+        install-torch: [0]
         # Groups are balanced on measured cold-compile-cache time, which is what a fresh runner pays.
         # tests/test_workflow_groups.py fails if a test file is missing from every group or listed twice.
         test-subset:
           - name: attention
             paths: |
               tests/test_attention.py
-              tests/dispatch/
           - name: layers
             paths: |
               tests/test_transformer.py
               tests/test_layers.py
+              tests/test_recurrent.py
+              tests/test_conv.py
+              tests/test_conv_transpose.py
+              tests/test_padding.py
+              tests/test_pool.py
               tests/rewriting/
           - name: optim
             paths: |
@@ -51,7 +60,27 @@ jobs:
               tests/test_workflow_groups.py
         # Windows runners are roughly twice as slow, and almost everything here is platform independent.
         # It runs one job over the parts that are not: file IO, paths, and compiling a graph to train.
+        # One job per backend, each installing only that backend. They live here rather than in
+        # `test-subset` because a matrix entry cannot change `os`, and mlx needs macOS.
         include:
+          - os: ubuntu-latest
+            python-version: "3.12"
+            install-jax: 1
+            test-subset:
+              name: dispatch-jax
+              paths: tests/dispatch/jax/
+          - os: macos-15
+            python-version: "3.12"
+            install-mlx: 1
+            test-subset:
+              name: dispatch-mlx
+              paths: tests/dispatch/mlx/
+          - os: ubuntu-latest
+            python-version: "3.12"
+            install-torch: 1
+            test-subset:
+              name: dispatch-torch
+              paths: tests/dispatch/pytorch/
           - os: windows-latest
             python-version: "3.12"
             test-subset:
@@ -71,8 +100,8 @@ jobs:
         shell: bash -l {0}
 
     steps:
-    - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
-    - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
       env:
         CACHE_NUMBER: 0
       with:
@@ -80,7 +109,7 @@ jobs:
         key: ${{ runner.os }}-py${{matrix.python-version}}-conda-${{ env.CACHE_NUMBER }}-${{
           hashFiles('conda_envs/pytensor_ml.yml') }}
     - name: Cache multiple paths
-      uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+      uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
       env:
         # Increase this value to reset cache if pyproject.toml has not changed
         CACHE_NUMBER: 0
@@ -92,7 +121,7 @@ jobs:
         key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{ env.CACHE_NUMBER }}-${{
           hashFiles('pyproject.toml') }}
     - name: Cache the pytensor compile directory
-      uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+      uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
       env:
         CACHE_NUMBER: 0
       with:
@@ -115,6 +144,17 @@ jobs:
         # IMPORTANT: This may break caching of conda packages! See https://github.com/conda-incubator/setup-miniconda/issues/267
         use-only-tar-bz2: false
 
+    - name: Install the backend this job tests
+      run: |
+        conda activate pytensor_ml
+        if [[ $INSTALL_JAX == "1" ]]; then pip install "jax>=0.8,<0.9.1" jaxlib; fi
+        if [[ $INSTALL_MLX == "1" ]]; then pip install "mlx>=0.30,<0.32"; fi
+        if [[ $INSTALL_TORCH == "1" ]]; then pip install torch --index-url https://download.pytorch.org/whl/cpu; fi
+      env:
+        INSTALL_JAX: ${{ matrix.install-jax }}
+        INSTALL_MLX: ${{ matrix.install-mlx }}
+        INSTALL_TORCH: ${{ matrix.install-torch }}
+
     - name: Install current branch
       run: |
         conda activate pytensor_ml
@@ -125,7 +165,7 @@ jobs:
       run: |
         python -m pytest -vv --durations=25 --cov=pytensor_ml --cov-report=xml --no-cov-on-fail --cov-report term $TEST_SUBSET
     - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1
+      uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0
       with:
         token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
         env_vars: TEST_SUBSET

pytensor_ml-0.2.0/PKG-INFO

@@ -0,0 +1,43 @@
+Metadata-Version: 2.5
+Name: pytensor_ml
+Version: 0.2.0
+Summary: A(nother) deep learning library, built on top of Pytensor
+Project-URL: Homepage, https://github.com/pymc-devs/pytensor_ml
+Project-URL: Source, https://github.com/pymc-devs/pytensor_ml
+Project-URL: Tracker, https://github.com/pymc-devs/pytensor_ml/issues
+Author: Ricardo Vieira
+Author-email: Jesse Grabowski <jessegrabowski@gmail.com>
+License-Expression: Apache-2.0
+License-File: LICENSE
+Keywords: autograd,deep learning,gradient descent,neural network,pymc
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: Unix
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.12
+Requires-Dist: numpy
+Requires-Dist: pytensor<4.0.0,>=3.2.3
+Requires-Dist: safetensors
+Provides-Extra: dev
+Requires-Dist: jax; extra == 'dev'
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pyyaml; extra == 'dev'
+Requires-Dist: scikit-learn; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: matplotlib; extra == 'examples'
+Requires-Dist: tqdm; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# pytensor_ml
+Neural network package built on Pytensor

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/conda_envs/pytensor_ml.yml

@@ -5,7 +5,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.2.3,<4.0.0
+  - pytensor>=3.3.0,<3.4.0
   - numpy
   - safetensors
   - scikit-learn
@@ -29,9 +29,10 @@ dependencies:
 
   # developer tools
   - pre-commit
-  - mypy
+  # Pinned rather than floated: a local run that disagrees with CI about the version is a local run
+  # that cannot predict CI, and both of these have shipped errors the other version does not report.
+  - mypy>=2.3.1,<2.4.0
 
-  # JAX backend for dispatch tests; via pip so Windows pulls CPU wheels instead of conda jaxlib
+  # No backend goes here. Each backend's dispatch tests get their own CI job that installs just that
+  # one, so the core jobs stay light and prove the library works with no backend present at all.
   - pip
-  - pip:
-      - jax

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/examples/mnist_feed_forward.ipynb

@@ -37,7 +37,7 @@
     "from sklearn.preprocessing import MinMaxScaler, OneHotEncoder\n",
     "\n",
     "from pytensor_ml.activations import LeakyReLU\n",
-    "from pytensor_ml.layers import BatchNorm2D, Dropout, Input, Linear, Sequential\n",
+    "from pytensor_ml.layers import BatchNorm, Dropout, Input, Linear, Sequential\n",
     "from pytensor_ml.loss import CrossEntropy\n",
     "from pytensor_ml.model import Model\n",
     "from pytensor_ml.optim import adagrad\n",
@@ -75,11 +75,11 @@
    "source": [
     "prediction_network = Sequential(\n",
     "    Linear(\"Linear_1\", n_in=64, n_out=256),\n",
-    "    BatchNorm2D(\"BatchNorm_1\"),\n",
+    "    BatchNorm(\"BatchNorm_1\"),\n",
     "    Dropout(\"Dropout_1\", p=0.5),\n",
     "    LeakyReLU(),\n",
     "    Linear(\"Linear_2\", n_in=256, n_out=128),\n",
-    "    BatchNorm2D(\"BatchNorm_2\"),\n",
+    "    BatchNorm(\"BatchNorm_2\"),\n",
     "    Dropout(\"Dropout_2\", p=0.5),\n",
     "    LeakyReLU(),\n",
     "    Linear(\"Logits\", n_in=128, n_out=10),\n",

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pyproject.toml

@@ -10,7 +10,8 @@ description = "A(nother) deep learning library, built on top of Pytensor"
 readme = "README.md"
 requires-python = ">=3.12"
 authors = [{name="Jesse Grabowski", email='jessegrabowski@gmail.com'}, {name="Ricardo Vieira"}]
-license = { file = "LICENSE" }
+license = "Apache-2.0"
+license-files = ["LICENSE"]
 classifiers = [
   "Development Status :: 2 - Pre-Alpha",
   "Intended Audience :: Science/Research",
@@ -76,6 +77,13 @@ filterwarnings = [
     # warnings filter per-test and undoes it. Numba emits it from the numba
     # linker, which is pytensor 3's default.
     "ignore::numba.core.errors.NumbaPerformanceWarning",
+    # torch.compile pulls in inductor, which imports torch.utils.mkldnn, whose class bodies still call
+    # the deprecated torch.jit.script_method. Raising here aborts the import half-done, so every later
+    # torch-mode test retries it and fails the same way. The message past this prefix depends on the
+    # Python version.
+    "ignore:`torch.jit.script_method` is:DeprecationWarning",
+    # pytensor hands torch its constants as the read-only numpy arrays it stores them in.
+    "ignore:The given NumPy array is not writable:UserWarning",
 ]
 
 [tool.ruff]

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/__init__.py

@@ -2,6 +2,7 @@ import logging
 import sys
 
 import pytensor_ml.dispatch  # installs the lazy backend-dispatch registration hook
+import pytensor_ml.rewriting.conv  # registers the convolution rewrites into pytensor's databases
 
 from pytensor_ml._version import __version__
 from pytensor_ml.checkpoint import load_state, save_state

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.0'
-__version_tuple__ = version_tuple = (0, 1, 0)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/base.py

@@ -7,6 +7,15 @@ from pytensor.compile.builders import SymbolicOp
 from pytensor.tensor.variable import TensorVariable
 
 
+def _check_input_rank(X: TensorVariable, name: str, n_spatial: int) -> None:
+    """Reject an input whose rank is not batch, one axis per spatial dimension, then channels."""
+    if X.ndim != n_spatial + 2:
+        raise ValueError(
+            f"{name} takes an input of shape (batch, {', '.join(['spatial'] * n_spatial)}, channels), "
+            f"so it needs a {n_spatial + 2}-dimensional input; got a {X.ndim}-dimensional one."
+        )
+
+
 class Layer(ABC):
     """Base class for the objects that build layer graphs. Defined here, not in ``pytensor_ml.layers``, so
     that ``pytensor_ml.activations`` can subclass it without a circular import."""

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/dispatch/__init__.py

@@ -9,6 +9,8 @@ import sys
 _REGISTRATIONS = {
     "pytensor.link.jax.dispatch": "pytensor_ml.dispatch.jax",
     "pytensor.link.mlx.dispatch": "pytensor_ml.dispatch.mlx",
+    "pytensor.link.numba.dispatch": "pytensor_ml.dispatch.numba",
+    "pytensor.link.pytorch.dispatch": "pytensor_ml.dispatch.pytorch",
 }
 
 

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/dispatch/jax/__init__.py

@@ -1,3 +1,5 @@
 # Imported for its side effect: registering pytensor_ml's JAX funcify dispatches. One submodule per
 # marker op that gets a kernel, mirroring the layout under pytensor_ml/layers.
 import pytensor_ml.dispatch.jax.attention
+import pytensor_ml.dispatch.jax.conv
+import pytensor_ml.dispatch.jax.pooling

pytensor_ml-0.2.0/pytensor_ml/dispatch/jax/conv.py

@@ -0,0 +1,72 @@
+import jax
+
+from pytensor.link.jax.dispatch import jax_funcify
+
+from pytensor_ml.layers.conv import ConvLayer, ConvLayerGrad
+
+
+def _spatial_letters(n_spatial: int) -> str:
+    """Name each spatial axis with a distinct letter, which is how jax takes a convolution's layout."""
+    letters = "XYZUVW"
+    if n_spatial > len(letters):
+        raise ValueError(
+            f"A convolution over {n_spatial} spatial axes has more axes than there are letters to name "
+            f"them with; {len(letters)} is the most jax supports here."
+        )
+    return letters[:n_spatial]
+
+
+def _convolution(op):
+    """The convolution both dispatches here run, so the gradient differentiates what the forward does."""
+    spatial = _spatial_letters(len(op.kernel_size))
+    # Our layouts are already the ones jax names: the input is channels-last and the kernel is taps,
+    # then input channels, then output. Nothing is transposed on the way in or out.
+    dimension_numbers = (f"N{spatial}C", f"{spatial}IO", f"N{spatial}C")
+
+    def convolve(X, W):
+        # Padding is a `pt.pad` node ahead of the op, so what arrives here is already padded.
+        return jax.lax.conv_general_dilated(
+            X,
+            W,
+            window_strides=op.stride,
+            padding="VALID",
+            rhs_dilation=op.dilation,
+            dimension_numbers=dimension_numbers,
+        )
+
+    return convolve
+
+
+@jax_funcify.register(ConvLayer)
+def jax_funcify_ConvLayer(op, node=None, **kwargs):
+    """Dispatch the convolution marker to ``jax.lax.conv_general_dilated`` (XLA/cuDNN)."""
+    convolve = _convolution(op)
+
+    def conv(X, W, *bias):
+        out = convolve(X, W)
+        return out + bias[0] if bias else out
+
+    return conv
+
+
+@jax_funcify.register(ConvLayerGrad)
+def jax_funcify_ConvLayerGrad(op, node=None, **kwargs):
+    """Let jax differentiate its own convolution rather than spelling the two gradients out."""
+    convolve = _convolution(op)
+    compute_dX, compute_dW = op.compute_dX, op.compute_dW
+
+    def conv_grad(X, W, cotangent):
+        # The vjp is taken only over the inputs whose gradient is wanted; the rest are closed over as
+        # constants, so jax never differentiates toward a gradient the graph will not read.
+        if compute_dX and compute_dW:
+            _, pullback = jax.vjp(convolve, X, W)
+            return tuple(pullback(cotangent))
+        # An op with one output is dispatched to a function returning that output, not a list of one.
+        if compute_dX:
+            _, pullback = jax.vjp(lambda x: convolve(x, W), X)
+        else:
+            _, pullback = jax.vjp(lambda w: convolve(X, w), W)
+        (gradient,) = pullback(cotangent)
+        return gradient
+
+    return conv_grad

pytensor_ml-0.2.0/pytensor_ml/dispatch/jax/pooling.py

@@ -0,0 +1,60 @@
+import math
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+
+from pytensor_ml.layers.conv import PoolLayer, PoolLayerGrad
+
+
+def _pooling(op):
+    """The pooling both dispatches here run, so the gradient differentiates what the forward does."""
+    # Windows span the spatial axes only; batch and channels take a width of one and are never reduced.
+    window = (1, *op.kernel_size, 1)
+    strides = (1, *op.stride, 1)
+    dilation = (1, *op.dilation, 1)
+    taps = math.prod(op.kernel_size)
+
+    # Padding is a `pt.pad` node ahead of the op, so what arrives here is already padded.
+    reduce_window = partial(
+        jax.lax.reduce_window,
+        window_dimensions=window,
+        window_strides=strides,
+        padding="VALID",
+        window_dilation=dilation,
+    )
+
+    if op.reduction == "max":
+
+        def pool(X):
+            return reduce_window(X, -jnp.inf, jax.lax.max)
+
+    else:
+
+        def pool(X):
+            return reduce_window(X, 0.0, jax.lax.add) / taps
+
+    return pool
+
+
+@jax_funcify.register(PoolLayer)
+def jax_funcify_PoolLayer(op, node=None, **kwargs):
+    """Dispatch the pooling marker to ``jax.lax.reduce_window``."""
+    return _pooling(op)
+
+
+@jax_funcify.register(PoolLayerGrad)
+def jax_funcify_PoolLayerGrad(op, node=None, **kwargs):
+    """Let jax differentiate its own pooling rather than spelling the routing out."""
+    pool = _pooling(op)
+
+    def pool_grad(X, cotangent):
+        # An op with one output is dispatched to a function returning that output, not a list of one.
+        _, pullback = jax.vjp(pool, X)
+        (gradient,) = pullback(cotangent)
+        return gradient
+
+    return pool_grad

{pytensor_ml-0.1.0 → pytensor_ml-0.2.0}/pytensor_ml/dispatch/mlx/__init__.py

@@ -1,3 +1,5 @@
 # Imported for its side effect: registering pytensor_ml's MLX funcify dispatches. One submodule per
 # marker op that gets a kernel, mirroring the layout under pytensor_ml/layers.
 import pytensor_ml.dispatch.mlx.attention
+import pytensor_ml.dispatch.mlx.conv
+import pytensor_ml.dispatch.mlx.pooling

pytensor_ml-0.2.0/pytensor_ml/dispatch/mlx/conv.py

@@ -0,0 +1,62 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch import mlx_funcify
+
+from pytensor_ml.layers.conv import ConvLayer, ConvLayerGrad
+
+
+def _convolution(op):
+    """The convolution both dispatches here run, so the gradient differentiates what the forward does."""
+    no_padding = (0,) * len(op.kernel_size)
+
+    def convolve(X, W):
+        # mlx takes the input channels-last as we store it, but wants the kernel's output channels
+        # leading rather than trailing. Padding is a `pt.pad` node ahead of the op.
+        return mx.conv_general(
+            X,
+            mx.moveaxis(W, -1, 0),
+            stride=op.stride,
+            padding=no_padding,
+            kernel_dilation=op.dilation,
+        )
+
+    return convolve
+
+
+@mlx_funcify.register(ConvLayer)
+def mlx_funcify_ConvLayer(op, node=None, **kwargs):
+    """Dispatch the convolution marker to ``mx.conv_general`` (fused Metal kernel)."""
+    convolve = _convolution(op)
+
+    def conv(X, W, *bias):
+        out = convolve(X, W)
+        return out + bias[0] if bias else out
+
+    return conv
+
+
+@mlx_funcify.register(ConvLayerGrad)
+def mlx_funcify_ConvLayerGrad(op, node=None, **kwargs):
+    """Let mlx differentiate its own convolution rather than spelling the two gradients out.
+
+    The forward value the pullback is built from is never evaluated, and mlx is lazy, so it costs
+    nothing beyond the intermediates the gradients need anyway.
+    """
+    convolve = _convolution(op)
+    compute_dX, compute_dW = op.compute_dX, op.compute_dW
+
+    def conv_grad(X, W, cotangent):
+        # The vjp is taken only over the inputs whose gradient is wanted; the rest are closed over as
+        # constants, so mlx never differentiates toward a gradient the graph will not read.
+        if compute_dX and compute_dW:
+            _, pullback = mx.vjp(convolve, [X, W], [cotangent])
+            return tuple(pullback)
+        # An op with one output is dispatched to a function returning that output, not a list of one.
+        if compute_dX:
+            _, pullback = mx.vjp(lambda x: convolve(x, W), [X], [cotangent])
+        else:
+            _, pullback = mx.vjp(lambda w: convolve(X, w), [W], [cotangent])
+        (gradient,) = pullback
+        return gradient
+
+    return conv_grad

pytensor_ml-0.2.0/pytensor_ml/dispatch/mlx/pooling.py

@@ -0,0 +1,57 @@
+import mlx.core as mx
+import mlx.nn as mnn
+
+from pytensor.link.mlx.dispatch import mlx_funcify
+
+from pytensor_ml.layers.conv import PoolLayer, PoolLayerGrad
+
+_POOLS = {
+    ("max", 1): mnn.MaxPool1d,
+    ("max", 2): mnn.MaxPool2d,
+    ("max", 3): mnn.MaxPool3d,
+    ("mean", 1): mnn.AvgPool1d,
+    ("mean", 2): mnn.AvgPool2d,
+    ("mean", 3): mnn.AvgPool3d,
+}
+
+
+def _pooling(op):
+    """The pooling both dispatches here run, so the gradient differentiates what the forward does."""
+    n_spatial = len(op.kernel_size)
+    if any(spacing != 1 for spacing in op.dilation):
+        raise NotImplementedError(
+            f"mlx pools over adjacent positions only, so it cannot take the dilation {op.dilation}. "
+            "Pool without dilation, or run this graph on jax or the default backend."
+        )
+    if (op.reduction, n_spatial) not in _POOLS:
+        raise NotImplementedError(
+            f"mlx has no pooling over {n_spatial} spatial axes; it goes up to three."
+        )
+
+    # mlx pools channels-last as we store it, and takes no padding here because a `pt.pad` node ahead
+    # of the op has already applied it.
+    return _POOLS[(op.reduction, n_spatial)](kernel_size=op.kernel_size, stride=op.stride)
+
+
+@mlx_funcify.register(PoolLayer)
+def mlx_funcify_PoolLayer(op, node=None, **kwargs):
+    """Dispatch the pooling marker to mlx's own pooling (a fused Metal kernel)."""
+    pool = _pooling(op)
+
+    def pooled(X):
+        return pool(X)
+
+    return pooled
+
+
+@mlx_funcify.register(PoolLayerGrad)
+def mlx_funcify_PoolLayerGrad(op, node=None, **kwargs):
+    """Let mlx differentiate its own pooling rather than spelling the routing out."""
+    pool = _pooling(op)
+
+    def pool_grad(X, cotangent):
+        # An op with one output is dispatched to a function returning that output, not a list of one.
+        _, pullback = mx.vjp(pool, [X], [cotangent])
+        return pullback[0]
+
+    return pool_grad

pytensor_ml-0.2.0/pytensor_ml/dispatch/numba/__init__.py

@@ -0,0 +1,3 @@
+# Imported for its side effect: registering pytensor_ml's numba funcify dispatches. One submodule per
+# op that gets a kernel, mirroring the layout under pytensor_ml/layers.
+import pytensor_ml.dispatch.numba.conv