pytensor-ml 0.3.0__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.3.0 → 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.3.0 → 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,6 +32,7 @@ jobs:
             paths: |
               tests/test_transformer.py
               tests/test_layers.py
+              tests/test_positional.py
           - name: recurrent
             paths: |
               tests/test_recurrent.py
@@ -72,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: |
@@ -103,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
-      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
+    - uses: prefix-dev/setup-pixi@9dabb60412d3d2d967a8d682f7ce01317af40514 # v0.11.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.32.2,<0.33"; 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.3.0 → 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.3.0 → pytensor_ml-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pytensor_ml
-Version: 0.3.0
+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,7 +25,7 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering
 Requires-Python: >=3.12
 Requires-Dist: numpy
-Requires-Dist: pytensor<4.0.0,>=3.3.3
+Requires-Dist: pytensor<4.0.0,>=3.3.8
 Provides-Extra: dev
 Requires-Dist: cloudpickle; extra == 'dev'
 Requires-Dist: jax; extra == 'dev'

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.3.0 → pytensor_ml-0.4.0}/docs/source/api/layers.rst

@@ -72,6 +72,7 @@ Normalization and regularization
 
     BatchNorm
     LayerNorm
+    RMSNorm
     GroupNorm
     Dropout
 
@@ -102,3 +103,5 @@ Attention and transformers
     FeedForward
     TransformerBlock
     scaled_dot_product_attention
+    RotaryEmbedding
+    rotary_embedding

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

@@ -10,6 +10,7 @@ Training
     :toctree: generated/
 
     compile_train
+    build_train_updates
 
 Update rules
 ------------

{pytensor_ml-0.3.0 → 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.3.0 → pytensor_ml-0.4.0}/docs/source/references.bib

@@ -71,3 +71,17 @@
   pages = {170--192},
   year = {2005},
 }
+
+@article{su2021roformer,
+  title = {RoFormer: Enhanced Transformer with Rotary Position Embedding},
+  author = {Su, Jianlin and Lu, Yu and Pan, Shengfeng and Murtadha, Ahmed and Wen, Bo and Liu, Yunfeng},
+  journal = {arXiv preprint arXiv:2104.09864},
+  year = {2021},
+}
+
+@article{chen2023interpolation,
+  title = {Extending Context Window of Large Language Models via Positional Interpolation},
+  author = {Chen, Shouyuan and Wong, Sherman and Chen, Liangjian and Tian, Yuandong},
+  journal = {arXiv preprint arXiv:2306.15595},
+  year = {2023},
+}