pytensor-ml 0.2.2__tar.gz → 0.4.0__tar.gz

pytensor_ml-0.4.0/.github/workflows/mypy.yml

@@ -0,0 +1,34 @@
+name: mypy
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+    paths:
+      - ".github/workflows/mypy.yml"
+      - "pytensor_ml/**"
+      - "scripts/run_mypy.py"
+      - "pyproject.toml"
+      - "pixi.lock"
+
+# Cancels all previous workflow runs for pull requests that have not completed.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  mypy:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - uses: prefix-dev/setup-pixi@9dabb60412d3d2d967a8d682f7ce01317af40514 # v0.11.0
+      with:
+        pixi-version: v0.81.0
+        environments: mypy
+        activate-environment: true
+        locked: true
+        cache: true
+
+    - name: Run mypy
+      run: python ./scripts/run_mypy.py

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/.github/workflows/pre-commit.yml

@@ -13,4 +13,9 @@ jobs:
     steps:
     - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
     - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+    # The pixi hooks call pixi directly and need no environment installed.
+    - uses: prefix-dev/setup-pixi@9dabb60412d3d2d967a8d682f7ce01317af40514 # v0.11.0
+      with:
+        pixi-version: v0.81.0
+        run-install: false
     - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/.github/workflows/run_tests.yml

@@ -14,17 +14,14 @@ concurrency:
 
 jobs:
   unittest:
-    name: unittest (${{ matrix.os }}, ${{ matrix.python-version }}, ${{ matrix.test-subset.name }})
+    name: unittest (${{ matrix.os }}, ${{ 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]
+        # The per-backend jobs under `include` swap in an environment carrying their one backend. A core
+        # job's environment has no backend at all, which is what proves the library runs without one.
+        pixi-environment: [test]
         # 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:
@@ -35,15 +32,25 @@ jobs:
             paths: |
               tests/test_transformer.py
               tests/test_layers.py
+              tests/test_positional.py
+          - name: recurrent
+            paths: |
               tests/test_recurrent.py
+              tests/rewriting/
+          - name: convolution
+            paths: |
               tests/test_conv.py
               tests/test_conv_transpose.py
               tests/test_padding.py
               tests/test_pool.py
-              tests/rewriting/
+              tests/test_upsample.py
           - name: optim
             paths: |
               tests/optim/
+          - name: models
+            paths: |
+              tests/test_pretrained.py
+              tests/models/
           - name: rest
             paths: |
               tests/test_serialize.py
@@ -51,7 +58,6 @@ jobs:
               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
@@ -64,25 +70,22 @@ jobs:
         # `test-subset` because a matrix entry cannot change `os`, and mlx needs macOS.
         include:
           - os: ubuntu-latest
-            python-version: "3.12"
-            install-jax: 1
+            pixi-environment: test-jax
             test-subset:
               name: dispatch-jax
               paths: tests/dispatch/jax/
           - os: macos-15
-            python-version: "3.12"
-            install-mlx: 1
+            pixi-environment: test-mlx
             test-subset:
               name: dispatch-mlx
               paths: tests/dispatch/mlx/
           - os: ubuntu-latest
-            python-version: "3.12"
-            install-torch: 1
+            pixi-environment: test-torch
             test-subset:
               name: dispatch-torch
               paths: tests/dispatch/pytorch/
           - os: windows-latest
-            python-version: "3.12"
+            pixi-environment: test
             test-subset:
               name: smoke
               paths: |
@@ -95,73 +98,33 @@ jobs:
     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
+    - uses: prefix-dev/setup-pixi@9dabb60412d3d2d967a8d682f7ce01317af40514 # v0.11.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') }}
+        pixi-version: v0.81.0
+        environments: ${{ matrix.pixi-environment }}
+        activate-environment: true
+        locked: true
+        cache: true
+    # A saved cache entry is immutable, so a key without the commit would freeze the cache at the first
+    # run on each lockfile, and every graph changed after it would compile cold on every run. Keying on
+    # the commit saves a fresh entry each run, and the restore keys pick up the newest one.
     - 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') }}
+        # pytensor compiles into %LOCALAPPDATA%\PyTensor on Windows and ~/.pytensor elsewhere.
+        path: ${{ runner.os == 'Windows' && '~/AppData/Local/PyTensor' || '~/.pytensor' }}
+        key: ${{ runner.os }}-pytensor-${{ env.CACHE_NUMBER }}-${{
+          matrix.test-subset.name }}-${{ hashFiles('pixi.lock') }}-${{ github.sha }}
         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
+          ${{ runner.os }}-pytensor-${{ env.CACHE_NUMBER }}-${{ matrix.test-subset.name }}-${{ hashFiles('pixi.lock') }}-
+          ${{ runner.os }}-pytensor-${{ env.CACHE_NUMBER }}-${{ matrix.test-subset.name }}-
 
     - name: Run tests
+      shell: bash
       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

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/.pre-commit-config.yaml

@@ -32,6 +32,31 @@ repos:
 
 - repo: local
   hooks:
+    - id: pixi-lock-check
+      name: pixi.lock is in sync with pyproject.toml
+      entry: pixi lock --check
+      language: system
+      files: ^(pyproject\.toml|pixi\.lock)$
+      pass_filenames: false
+
+    - id: pixi-export-dev
+      name: conda_envs/pytensor_ml.yml is exported from the pixi dev environment
+      # Called directly rather than through a pixi task so the hook reads the manifest only; `pixi run`
+      # would re-solve the lock file and install the environment first. --no-pypi drops the editable
+      # install, because conda runs pip from the env file's directory, where `-e .` names conda_envs/.
+      # -p pins the export to one platform, so it does not depend on the machine running the hook.
+      entry: pixi workspace export conda-environment -e dev -n pytensor_ml -p linux-64 --no-pypi conda_envs/pytensor_ml.yml
+      language: system
+      files: ^(pyproject\.toml|pixi\.lock|conda_envs/pytensor_ml\.yml)$
+      pass_filenames: false
+
+    - id: pixi-export-docs
+      name: conda_envs/environment-docs.yml is exported from the pixi docs environment
+      entry: pixi workspace export conda-environment -e docs -n pytensor_ml-docs -p linux-64 --no-pypi conda_envs/environment-docs.yml
+      language: system
+      files: ^(pyproject\.toml|pixi\.lock|conda_envs/environment-docs\.yml)$
+      pass_filenames: false
+
     - id: notebook_name
       entry: '\(notebook_name\)='
       language: pygrep

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pytensor_ml
-Version: 0.2.2
+Version: 0.4.0
 Summary: A(nother) deep learning library, built on top of Pytensor
 Project-URL: Homepage, https://github.com/pymc-devs/pytensor-ml
 Project-URL: Documentation, https://pytensor-ml.readthedocs.io
@@ -25,15 +25,16 @@ 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
+Requires-Dist: pytensor<4.0.0,>=3.3.8
 Provides-Extra: dev
+Requires-Dist: cloudpickle; 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: safetensors; extra == 'dev'
 Requires-Dist: scikit-learn; extra == 'dev'
 Provides-Extra: examples
 Requires-Dist: matplotlib; extra == 'examples'
@@ -61,8 +62,9 @@ any other PyTensor graph — a PyMC model included — as there is nothing else
 pip install pytensor-ml
 ```
 
-The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
-`jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
+The only hard dependencies are `pytensor` and `numpy`. Saving and loading weights needs `safetensors`
+(`pip install safetensors`). A backend beyond the default (`numba`, `jax`, `torch`, `mlx`) is installed
+separately, and only loads when you actually compile against it.
 
 ## Quickstart
 
@@ -93,11 +95,11 @@ network = Sequential(
     ReLU(),
     Linear("logits", n_in=128, n_out=10),
 )
-model = Model(X_in, network(X_in)).initialize(seed=0)
+model = Model(network(X_in)).initialize(seed=0)
 
 rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
 loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-step = model.compile_train(rule, loss_fn, ndim_out=2)
+step = model.compile_train(rule, loss_fn)
 
 loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
 for _ in range(500):

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/README.md

@@ -19,8 +19,9 @@ any other PyTensor graph — a PyMC model included — as there is nothing else
 pip install pytensor-ml
 ```
 
-The only hard dependencies are `pytensor`, `numpy`, and `safetensors`. A backend beyond the default (`numba`,
-`jax`, `torch`, `mlx`) is installed separately, and only loads when you actually compile against it.
+The only hard dependencies are `pytensor` and `numpy`. Saving and loading weights needs `safetensors`
+(`pip install safetensors`). A backend beyond the default (`numba`, `jax`, `torch`, `mlx`) is installed
+separately, and only loads when you actually compile against it.
 
 ## Quickstart
 
@@ -51,11 +52,11 @@ network = Sequential(
     ReLU(),
     Linear("logits", n_in=128, n_out=10),
 )
-model = Model(X_in, network(X_in)).initialize(seed=0)
+model = Model(network(X_in)).initialize(seed=0)
 
 rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
 loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-step = model.compile_train(rule, loss_fn, ndim_out=2)
+step = model.compile_train(rule, loss_fn)
 
 loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
 for _ in range(500):

pytensor_ml-0.4.0/conda_envs/environment-docs.yml

@@ -0,0 +1,24 @@
+name: pytensor_ml-docs
+channels:
+- conda-forge
+- nodefaults
+dependencies:
+- safetensors *
+- matplotlib *
+- ipython *
+- jupyter *
+- sphinx >=7
+- pydata-sphinx-theme *
+- myst-nb *
+- numpydoc *
+- sphinx-copybutton *
+- sphinx-design *
+- sphinx-codeautolink *
+- sphinx-sitemap *
+- sphinx-notfound-page *
+- sphinx-autobuild *
+- jupyter-sphinx *
+- sphinxcontrib-bibtex *
+- pytensor >=3.3.8,<3.4.0
+- numpy *
+- python >=3.12

pytensor_ml-0.4.0/conda_envs/pytensor_ml-gpu_jax.yml

@@ -0,0 +1,22 @@
+name: pytensor_ml-gpu_jax
+channels:
+- conda-forge
+- nodefaults
+dependencies:
+- safetensors *
+- scikit-learn *
+- cloudpickle *
+- pytest *
+- pytest-cov *
+- pytest-xdist *
+- pytest-benchmark *
+- pytest-mock *
+- pyyaml *
+- mypy >=2.4.0,<2.5.0
+- pre-commit *
+- jax >=0.12,<0.13
+- jaxlib >=0.12,<0.13 cuda*
+- cuda-nvcc *
+- pytensor >=3.3.8,<3.4.0
+- numpy *
+- python >=3.12

pytensor_ml-0.4.0/conda_envs/pytensor_ml-gpu_metal.yml

@@ -0,0 +1,20 @@
+name: pytensor_ml-gpu_metal
+channels:
+- conda-forge
+- nodefaults
+dependencies:
+- safetensors *
+- scikit-learn *
+- cloudpickle *
+- pytest *
+- pytest-cov *
+- pytest-xdist *
+- pytest-benchmark *
+- pytest-mock *
+- pyyaml *
+- mypy >=2.4.0,<2.5.0
+- pre-commit *
+- mlx >=0.32.4,<0.33
+- pytensor >=3.3.8,<3.4.0
+- numpy *
+- python >=3.12

pytensor_ml-0.4.0/conda_envs/pytensor_ml.yml

@@ -0,0 +1,19 @@
+name: pytensor_ml
+channels:
+- conda-forge
+- nodefaults
+dependencies:
+- safetensors *
+- scikit-learn *
+- cloudpickle *
+- pytest *
+- pytest-cov *
+- pytest-xdist *
+- pytest-benchmark *
+- pytest-mock *
+- pyyaml *
+- mypy >=2.4.0,<2.5.0
+- pre-commit *
+- pytensor >=3.3.8,<3.4.0
+- numpy *
+- python >=3.12

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/api/activations.rst

@@ -10,6 +10,7 @@ Activations
     ReLU
     LeakyReLU
     GELU
+    QuickGELU
     Swish
     Sigmoid
     SoftPlus

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/api/layers.rst

@@ -46,6 +46,8 @@ Convolution and pooling
     MaxPool2D
     AvgPool1D
     AvgPool2D
+    Upsample1D
+    Upsample2D
 
 Padding
 -------
@@ -70,6 +72,8 @@ Normalization and regularization
 
     BatchNorm
     LayerNorm
+    RMSNorm
+    GroupNorm
     Dropout
 
 Recurrent
@@ -99,3 +103,5 @@ Attention and transformers
     FeedForward
     TransformerBlock
     scaled_dot_product_attention
+    RotaryEmbedding
+    rotary_embedding

pytensor_ml-0.4.0/docs/source/api/models.rst

@@ -0,0 +1,34 @@
+Pretrained models
+=================
+
+.. currentmodule:: pytensor_ml.models
+
+Dispatch
+--------
+
+.. autosummary::
+    :toctree: generated/
+
+    architecture_name
+    build_from_config
+    register_builder
+
+Blocks
+------
+
+.. autosummary::
+    :toctree: generated/
+
+    ResnetBlock2D
+    AttentionBlock2D
+
+Binding weights
+---------------
+
+.. autosummary::
+    :toctree: generated/
+
+    KeyMap
+    bind_layer_norm
+    bind_linear
+    channels_last

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/api/optim.rst

@@ -10,6 +10,7 @@ Training
     :toctree: generated/
 
     compile_train
+    build_train_updates
 
 Update rules
 ------------
@@ -26,6 +27,16 @@ Update rules
     rprop
     adagrad
     adadelta
+    lbfgs
+
+Line searches
+-------------
+
+.. autosummary::
+    :toctree: generated/
+
+    zoom_line_search
+    LineSearch
 
 Transforms
 ----------
@@ -73,6 +84,7 @@ eliminated the value's last consumer.
 
     constant_schedule
     linear_schedule
+    linear_onecycle_schedule
     cosine_schedule
     exponential_schedule
     polynomial_schedule
@@ -85,9 +97,7 @@ Building blocks
 .. autosummary::
     :toctree: generated/
 
-    counter
     get_gradients
-    reuses_state
     scalar_state
     state_for
     steps_of
@@ -131,3 +141,11 @@ Low-level update functions
     rprop_updates
     adagrad_updates
     adadelta_updates
+    lbfgs_updates
+
+.. currentmodule:: pytensor_ml.optim.lbfgs
+
+.. autosummary::
+    :toctree: generated/
+
+    LBFGSDirection

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/api/state.rst

@@ -33,5 +33,7 @@ Initializers
     XavierUniformInitializer
     OrthogonalInitializer
     UnrecordedInitializer
+    EmptyInitializer
+    initial_values_from
     initialize_params
     fans

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/api.rst

@@ -14,5 +14,6 @@ API Reference
     api/optim
     api/state
     api/serialization
+    api/models
     api/pytensorf
     api/util

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/dev/contributing.rst

@@ -18,13 +18,23 @@ Setting up
     pip install -e ".[dev]"
     pre-commit install
 
-The conda environments under ``conda_envs/`` are what CI installs, pinned
-harder than the package metadata. Reach for one when a failure reproduces in CI
-but not locally:
+CI installs the pixi environments locked in ``pixi.lock``, pinned harder than
+the package metadata. Reach for one when a failure reproduces in CI but not
+locally:
+
+.. code-block:: bash
+
+    pixi run -e test pytest             # the core jobs; test-jax, test-mlx, test-torch for a backend
+    pixi run -e mypy python scripts/run_mypy.py
+
+``conda_envs/pytensor_ml.yml`` is exported from the pixi ``dev`` environment by a
+pre-commit hook, for setting up with conda instead:
 
 .. code-block:: bash
 
     conda env create -f conda_envs/pytensor_ml.yml
+    conda activate pytensor_ml
+    pip install -e .
 
 Docs are built with pixi, which owns its own environment:
 

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/dev/style_guide.rst

@@ -180,7 +180,7 @@ Every public entrypoint carries an ``Examples`` section. Three rules:
         from pytensor_ml.optim import adam, chain, clip_by_global_norm, compile_train
 
         X = Input("X", shape=(None, 4))
-        loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError(), ndim_out=2)
+        loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())
 
         step = compile_train(loss, chain(clip_by_global_norm(1.0), adam(1e-3)))
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))

{pytensor_ml-0.2.2 → pytensor_ml-0.4.0}/docs/source/index.rst

@@ -41,6 +41,9 @@ Quick example
 
     import numpy as np
 
+    from pytensor import config
+    from sklearn.datasets import load_digits
+
     from pytensor_ml.activations import ReLU
     from pytensor_ml.layers import Input, Linear, Sequential
     from pytensor_ml.loss import CrossEntropy
@@ -48,17 +51,21 @@ Quick example
     from pytensor_ml.optim import adam, chain, clip_by_global_norm, cosine_schedule
     from pytensor_ml.util import DataLoader
 
+    X, y = load_digits(return_X_y=True)
+    X = (X / 16.0).astype(config.floatX)
+    y_onehot = np.eye(10, dtype=config.floatX)[y]
+
     X_in = Input("X_in", shape=(None, 64))
     network = Sequential(
         Linear("fc1", n_in=64, n_out=128),
         ReLU(),
         Linear("logits", n_in=128, n_out=10),
     )
-    model = Model(X_in, network(X_in)).initialize(seed=0)
+    model = Model(network(X_in)).initialize(seed=0)
 
     rule = chain(clip_by_global_norm(1.0), adam(learning_rate=cosine_schedule(1e-3, total_steps=500)))
     loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
-    step = model.compile_train(rule, loss_fn, ndim_out=2)
+    step = model.compile_train(rule, loss_fn)
 
     loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
     for _ in range(500):