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

{pytensor_ml-0.0.7 → 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.0.7 → 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.0.7 → 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.2.0/.github/workflows/run_tests.yml

@@ -0,0 +1,174 @@
+name: run_tests
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+# Cancels all previous workflow runs for pull requests that have not completed.
+concurrency:
+  # The concurrency group contains the workflow name and the branch name for pull requests
+  # or the commit hash for any other events.
+  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  unittest:
+    name: unittest (${{ matrix.os }}, ${{ matrix.python-version }}, ${{ matrix.test-subset.name }})
+    strategy:
+      fail-fast: false
+      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
+          - 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: |
+              tests/optim/
+          - name: rest
+            paths: |
+              tests/test_serialize.py
+              tests/test_activations.py
+              tests/test_loss.py
+              tests/test_checkpoint.py
+              tests/test_model.py
+              tests/test_pretrained.py
+              tests/test_pytensorf.py
+              tests/test_state.py
+              tests/test_base.py
+              tests/test_collect.py
+              tests/test_util.py
+              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:
+              name: smoke
+              paths: |
+                tests/test_checkpoint.py
+                tests/test_serialize.py
+                tests/optim/test_training.py
+
+    runs-on: ${{ matrix.os }}
+
+    env:
+      TEST_SUBSET: ${{ matrix.test-subset.paths }}
+
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+      env:
+        CACHE_NUMBER: 0
+      with:
+        path: ~/conda_pkgs_dir
+        key: ${{ runner.os }}-py${{matrix.python-version}}-conda-${{ env.CACHE_NUMBER }}-${{
+          hashFiles('conda_envs/pytensor_ml.yml') }}
+    - name: Cache multiple paths
+      uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+      env:
+        # Increase this value to reset cache if pyproject.toml has not changed
+        CACHE_NUMBER: 0
+      with:
+        path: |
+          ~/.cache/pip
+          $RUNNER_TOOL_CACHE/Python/*
+          ~\AppData\Local\pip\Cache
+        key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{ env.CACHE_NUMBER }}-${{
+          hashFiles('pyproject.toml') }}
+    - name: Cache the pytensor compile directory
+      uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+      env:
+        CACHE_NUMBER: 0
+      with:
+        path: ~/.pytensor
+        key: ${{ runner.os }}-py${{ matrix.python-version }}-pytensor-${{ env.CACHE_NUMBER }}-${{
+          matrix.test-subset.name }}-${{ hashFiles('conda_envs/pytensor_ml.yml') }}
+        restore-keys: |
+          ${{ runner.os }}-py${{ matrix.python-version }}-pytensor-${{ env.CACHE_NUMBER }}-${{ matrix.test-subset.name }}-
+    - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1
+      with:
+        miniforge-variant: Miniforge3
+        miniforge-version: latest
+        mamba-version: "*"
+        activate-environment: pytensor_ml
+        channel-priority: strict
+        environment-file: conda_envs/pytensor_ml.yml
+        python-version: 3.12
+        use-mamba: true
+
+        # 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
+        pip install -e .
+        python --version
+
+    - name: Run tests
+      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@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
+        env_vars: TEST_SUBSET
+        name: ${{ matrix.os }}
+        fail_ci_if_error: false
+        verbose: true

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.2.0/conda_envs/pytensor_ml.yml

@@ -0,0 +1,38 @@
+name: pytensor_ml
+channels:
+  - conda-forge
+  - nodefaults
+
+dependencies:
+  - python>=3.12
+  - pytensor>=3.3.0,<3.4.0
+  - numpy
+  - safetensors
+  - scikit-learn
+
+  # For testing
+  - coveralls
+  - diff-cover
+  - pytest
+  - pytest-cov
+  - pytest-xdist
+  - pytest-benchmark
+  - pytest-mock
+  - pyyaml  # tests/test_workflow_groups.py reads the CI matrix
+
+  # For building docs
+  - sphinx
+  - sphinx_rtd_theme
+  - pygments
+  - pydot
+  - ipython
+
+  # developer tools
+  - pre-commit
+  # 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
+
+  # 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

{pytensor_ml-0.0.7 → 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, 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",
@@ -63,8 +63,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# X_in = Input('X_in', shape=(64, ))\n",
-    "X_in = pytensor.tensor.tensor(\"X_in\", shape=(None, 64))"
+    "X_in = Input(\"X_in\", shape=(None, 64))"
    ]
   },
   {
@@ -76,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.0.7 → 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",
@@ -52,6 +53,11 @@ dev = [
     "scikit-learn",
     "mypy",
     "jax",
+    "pyyaml",
+]
+examples = [
+    "matplotlib",
+    "tqdm",
 ]
 
 [tool.hatch.version]
@@ -71,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.0.7 → 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.0.7 → 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.0.7'
-__version_tuple__ = version_tuple = (0, 0, 7)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.0.7 → 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.0.7 → pytensor_ml-0.2.0}/pytensor_ml/checkpoint.py

@@ -42,19 +42,24 @@ def save_state(shared_variables: Sequence[SharedVariable], path: str | Path) ->
         Destination archive, written verbatim.
     """
     indexed = _index_by_name(shared_variables)
-    tensors = {name: _as_saveable_array(variable.get_value()) for name, variable in indexed.items()}
+    tensors = {
+        name: _as_saveable_array(variable.get_value(), variable.type.dtype)
+        for name, variable in indexed.items()
+    }
     save_file(tensors, fspath(path))
 
 
-def _as_saveable_array(value: Any) -> np.ndarray:
+def _as_saveable_array(value: Any, dtype: str) -> np.ndarray:
     """
-    Return a value as a C-contiguous numpy array safetensors will accept, preserving its rank.
+    Return a value as a C-contiguous numpy array of ``dtype`` that safetensors will accept, preserving its
+    rank.
 
     A JIT backend stores its own array type in the shared variables a compiled function updates, so a
     trained model's values are ``jax.Array`` or ``mlx.core.array``; both convert through the array
-    protocol.
+    protocol. It may also narrow them -- JAX computes in 32-bit unless ``floatX`` is ``float64`` -- so the
+    archive records the declared dtype to stay portable across precisions.
     """
-    array = np.asarray(value)
+    array = np.asarray(value, dtype=dtype)
     # Not np.ascontiguousarray: it forces ndim >= 1, reshaping a rank-0 array such as the step count.
     return array if array.flags["C_CONTIGUOUS"] else np.array(array, order="C")
 
@@ -110,11 +115,14 @@ def load_state(
     mismatches = []
     for key, variable in target_by_key.items():
         value = values[key]
-        # Through numpy: a backend array's own dtype object never compares equal to a numpy dtype.
+        # The declared dtype, not the stored value's, which a backend may have narrowed (see
+        # _as_saveable_array); set_value filters to the declaration anyway. Shape has to come from the
+        # value, whose runtime dimensions the declared type may leave unknown.
         current = np.asarray(variable.get_value(borrow=True))
-        if value.shape != current.shape or value.dtype != current.dtype:
+        declared_dtype = np.dtype(variable.type.dtype)
+        if value.shape != current.shape or value.dtype != declared_dtype:
             mismatches.append(
-                f"  {variable.name!r}: target is {current.dtype} of shape {current.shape}, "
+                f"  {variable.name!r}: target is {declared_dtype} of shape {current.shape}, "
                 f"archive has {value.dtype} of shape {value.shape}"
             )
     if mismatches:

{pytensor_ml-0.0.7 → 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.0.7 → 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.0.7 → 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