pytensor-ml 0.2.1__tar.gz → 0.3.0__tar.gz

pytensor_ml-0.3.0/.github/workflows/rtd-preview.yml

@@ -0,0 +1,30 @@
+name: Read the Docs PR preview
+
+on:
+  pull_request_target:
+    types:
+      - opened
+      - synchronize
+      - reopened
+    # Only fire when something that could change the docs build is touched.
+    paths:
+      - "docs/**"
+      - ".readthedocs.yaml"
+      - "conda_envs/environment-docs.yml"
+      - "examples/**"
+      - "pytensor_ml/**"
+      - "pyproject.toml"
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          # Project slug as configured on Read the Docs. Verify this matches the slug shown in the
+          # RTD project URL (https://readthedocs.org/projects/<slug>/) once the project is imported
+          # there.
+          project-slug: "pytensor-ml"

{pytensor_ml-0.2.1 → pytensor_ml-0.3.0}/.github/workflows/run_tests.yml

@@ -35,15 +35,24 @@ jobs:
             paths: |
               tests/test_transformer.py
               tests/test_layers.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 +60,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
@@ -148,7 +156,7 @@ jobs:
       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_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 }}

pytensor_ml-0.3.0/.readthedocs.yaml

@@ -0,0 +1,24 @@
+version: 2
+
+sphinx:
+  configuration: docs/source/conf.py
+  fail_on_warning: false
+
+conda:
+  environment: conda_envs/environment-docs.yml
+
+python:
+  install:
+    - method: pip
+      path: .
+
+build:
+  os: "ubuntu-22.04"
+  tools:
+    python: "miniforge3-latest"
+  jobs:
+    # hatch-vcs derives the version from git tags; RTD's default checkout is shallow and tagless, so
+    # fetch tags or pytensor_ml.__version__ falls back to 0.0.0+unknown and the version selector
+    # labels break.
+    post_checkout:
+      - git fetch --tags --unshallow || true

pytensor_ml-0.3.0/PKG-INFO

@@ -0,0 +1,140 @@
+Metadata-Version: 2.5
+Name: pytensor_ml
+Version: 0.3.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
+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.3.3
+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'
+Requires-Dist: tqdm; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# pytensor_ml
+
+A(nother) deep learning library, built on top of [PyTensor](https://github.com/pymc-devs/pytensor).
+
+Networks are ordinary PyTensor graphs. You build one out of layers, and everything PyTensor already does —
+symbolic differentiation, graph rewrites, and compilation to Numba, C, JAX, PyTorch, or MLX — applies to it
+unchanged. Training is a compiled function that takes a batch and returns a loss; there is no separate runtime
+or tape.
+
+That goes all the way down: layers are graph constructors, parameters are shared variables, and a training
+step is a compiled function whose updates are the optimizer. Because a model is only a graph, it composes with
+any other PyTensor graph — a PyMC model included — as there is nothing else to interoperate with.
+
+> **Status: pre-alpha.** The API is still moving, and there is no release-to-release compatibility guarantee yet.
+
+## Installation
+
+```bash
+pip install pytensor-ml
+```
+
+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
+
+Train a classifier on scikit-learn's digits, then run inference:
+
+```python
+import numpy as np
+import pytensor
+
+pytensor.config.floatX = "float32"
+
+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
+from pytensor_ml.model import Model
+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("float32")
+y_onehot = np.eye(10, dtype="float32")[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(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)
+
+loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
+for _ in range(500):
+    loss_value = step(*loader())
+
+accuracy = (model.predict(X).argmax(axis=-1) == y).mean()
+```
+
+`compile_train` builds the loss against a target placeholder, differentiates it, folds in any stateful layer
+updates (batch norm running statistics, RNG advances, the training clock a schedule reads), and compiles a
+one-step function. `predict` compiles a separate inference pass, with dropout removed and batch norm reading
+its running statistics.
+
+## Documentation
+
+The full API reference and user guide live at
+[pytensor-ml.readthedocs.io](https://pytensor-ml.readthedocs.io).
+
+For worked models end to end — training loops, convolutional and recurrent networks, transformers, saving and
+reloading — see the [examples gallery](https://pytensor-ml.readthedocs.io/en/latest/examples/gallery.html).
+
+## Contributing
+
+Contributions are welcome. To get set up:
+
+```bash
+pip install -e ".[dev]"
+pre-commit install
+pytest
+```
+
+Formatting and linting run through `ruff` under pre-commit, and `mypy` checks `pytensor_ml/`; both also run in
+CI. Bug reports and feature requests belong in the
+[issue tracker](https://github.com/pymc-devs/pytensor-ml/issues).
+
+## License
+
+Apache 2.0. See [LICENSE](LICENSE).

pytensor_ml-0.3.0/README.md

@@ -0,0 +1,97 @@
+# pytensor_ml
+
+A(nother) deep learning library, built on top of [PyTensor](https://github.com/pymc-devs/pytensor).
+
+Networks are ordinary PyTensor graphs. You build one out of layers, and everything PyTensor already does —
+symbolic differentiation, graph rewrites, and compilation to Numba, C, JAX, PyTorch, or MLX — applies to it
+unchanged. Training is a compiled function that takes a batch and returns a loss; there is no separate runtime
+or tape.
+
+That goes all the way down: layers are graph constructors, parameters are shared variables, and a training
+step is a compiled function whose updates are the optimizer. Because a model is only a graph, it composes with
+any other PyTensor graph — a PyMC model included — as there is nothing else to interoperate with.
+
+> **Status: pre-alpha.** The API is still moving, and there is no release-to-release compatibility guarantee yet.
+
+## Installation
+
+```bash
+pip install pytensor-ml
+```
+
+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
+
+Train a classifier on scikit-learn's digits, then run inference:
+
+```python
+import numpy as np
+import pytensor
+
+pytensor.config.floatX = "float32"
+
+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
+from pytensor_ml.model import Model
+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("float32")
+y_onehot = np.eye(10, dtype="float32")[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(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)
+
+loader = DataLoader(X, y_onehot, batch_size=64, random_state=0)
+for _ in range(500):
+    loss_value = step(*loader())
+
+accuracy = (model.predict(X).argmax(axis=-1) == y).mean()
+```
+
+`compile_train` builds the loss against a target placeholder, differentiates it, folds in any stateful layer
+updates (batch norm running statistics, RNG advances, the training clock a schedule reads), and compiles a
+one-step function. `predict` compiles a separate inference pass, with dropout removed and batch norm reading
+its running statistics.
+
+## Documentation
+
+The full API reference and user guide live at
+[pytensor-ml.readthedocs.io](https://pytensor-ml.readthedocs.io).
+
+For worked models end to end — training loops, convolutional and recurrent networks, transformers, saving and
+reloading — see the [examples gallery](https://pytensor-ml.readthedocs.io/en/latest/examples/gallery.html).
+
+## Contributing
+
+Contributions are welcome. To get set up:
+
+```bash
+pip install -e ".[dev]"
+pre-commit install
+pytest
+```
+
+Formatting and linting run through `ruff` under pre-commit, and `mypy` checks `pytensor_ml/`; both also run in
+CI. Bug reports and feature requests belong in the
+[issue tracker](https://github.com/pymc-devs/pytensor-ml/issues).
+
+## License
+
+Apache 2.0. See [LICENSE](LICENSE).

pytensor_ml-0.3.0/conda_envs/environment-docs.yml

@@ -0,0 +1,35 @@
+# The environment Read the Docs builds from. `pixi run docs-build` uses the equivalent list under
+# [tool.pixi.feature.docs] in pyproject.toml; the two have to move together.
+name: pytensor_ml-docs
+channels:
+  - conda-forge
+  - nodefaults
+
+dependencies:
+  - python>=3.12
+  # Runtime deps: autodoc imports pytensor_ml, so the full runtime stack has to be in scope.
+  - pytensor>=3.3.3,<3.4.0
+  - numpy
+  - safetensors
+  # The gallery extension renders notebook thumbnails with matplotlib.
+  - matplotlib
+
+  # Docs build deps.
+  - 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
+  - pip
+  - pip:
+      # pytensor_ml itself so autodoc resolves the current source tree.
+      - -e ..

{pytensor_ml-0.2.1 → pytensor_ml-0.3.0}/conda_envs/pytensor_ml-gpu_jax.yml

@@ -6,7 +6,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.2.3,<4.0.0
+  - pytensor>=3.3.3,<4.0.0
   - numpy
   - scikit-learn
 

{pytensor_ml-0.2.1 → pytensor_ml-0.3.0}/conda_envs/pytensor_ml.yml

@@ -5,12 +5,13 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.3.0,<3.4.0
+  - pytensor>=3.3.3,<3.4.0
   - numpy
   - safetensors
   - scikit-learn
 
   # For testing
+  - cloudpickle  # tests/optim/test_composition.py ships a configured rule to a worker
   - coveralls
   - diff-cover
   - pytest
@@ -20,13 +21,6 @@ dependencies:
   - 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

pytensor_ml-0.3.0/docs/.gitignore

@@ -0,0 +1,16 @@
+# Sphinx build output
+build/
+jupyter_execute/
+
+# Python cache
+__pycache__/
+*.pyc
+
+# Autosummary-generated stubs (one .rst per object, regenerated on build)
+source/api/generated/
+source/api/**/generated/
+source/api/**/classmethods/
+
+# Notebook gallery artifacts written by docs/sphinxext/generate_gallery.py
+source/examples/
+source/_thumbnails/

pytensor_ml-0.3.0/docs/Makefile

@@ -0,0 +1,34 @@
+# Minimal Makefile for Sphinx documentation.
+# Mirrors the standard sphinx-quickstart output.
+
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+.PHONY: help clean html show livehtml linkcheck Makefile
+
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+clean:
+	rm -rf "$(BUILDDIR)"
+	rm -rf "$(SOURCEDIR)"/api/generated
+	rm -rf "$(SOURCEDIR)"/api/*/generated
+	rm -rf "$(SOURCEDIR)"/_thumbnails
+	rm -f  "$(SOURCEDIR)"/examples/gallery.rst
+
+html:
+	@$(SPHINXBUILD) -M html "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+show: html
+	@open "$(BUILDDIR)/html/index.html"
+
+livehtml:
+	sphinx-autobuild "$(SOURCEDIR)" "$(BUILDDIR)/html" $(SPHINXOPTS) $(O)
+
+linkcheck:
+	@$(SPHINXBUILD) -M linkcheck "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

pytensor_ml-0.3.0/docs/make.bat

@@ -0,0 +1,29 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Install Sphinx, then re-run.
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

pytensor_ml-0.3.0/docs/source/_templates/autosummary/class.rst

@@ -0,0 +1,34 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+
+   {% block methods %}
+   {% if methods %}
+
+   .. rubric:: Methods
+
+   .. autosummary::
+      :toctree: classmethods
+
+   {% for item in methods %}
+   {%- if item not in inherited_members %}
+      {{ objname }}.{{ item }}
+   {% endif %}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: Attributes
+
+   .. autosummary::
+   {% for item in attributes %}
+   {%- if item not in inherited_members %}
+      ~{{ name }}.{{ item }}
+   {% endif %}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}

pytensor_ml-0.3.0/docs/source/api/activations.rst

@@ -0,0 +1,18 @@
+Activations
+===========
+
+.. currentmodule:: pytensor_ml.activations
+
+.. autosummary::
+    :toctree: generated/
+
+    Activation
+    ReLU
+    LeakyReLU
+    GELU
+    QuickGELU
+    Swish
+    Sigmoid
+    SoftPlus
+    Softmax
+    Tanh

pytensor_ml-0.3.0/docs/source/api/layers.rst

@@ -0,0 +1,104 @@
+Layers
+======
+
+.. currentmodule:: pytensor_ml.layers
+
+Base
+----
+
+.. autosummary::
+    :toctree: generated/
+
+    Layer
+
+Combinators
+-----------
+
+.. autosummary::
+    :toctree: generated/
+
+    Input
+    Sequential
+    Concatenate
+    Flatten
+    Squeeze
+
+Dense
+-----
+
+.. autosummary::
+    :toctree: generated/
+
+    Linear
+    Embedding
+
+Convolution and pooling
+-----------------------
+
+.. autosummary::
+    :toctree: generated/
+
+    Conv1D
+    Conv2D
+    ConvTranspose1D
+    ConvTranspose2D
+    MaxPool1D
+    MaxPool2D
+    AvgPool1D
+    AvgPool2D
+    Upsample1D
+    Upsample2D
+
+Padding
+-------
+
+.. autosummary::
+    :toctree: generated/
+
+    ZeroPad1D
+    ZeroPad2D
+    ConstantPad1D
+    ConstantPad2D
+    ReflectionPad1D
+    ReflectionPad2D
+    ReplicationPad1D
+    ReplicationPad2D
+
+Normalization and regularization
+--------------------------------
+
+.. autosummary::
+    :toctree: generated/
+
+    BatchNorm
+    LayerNorm
+    GroupNorm
+    Dropout
+
+Recurrent
+---------
+
+.. autosummary::
+    :toctree: generated/
+
+    RNN
+    LSTM
+    GRU
+    Bidirectional
+    Recurrent
+    RecurrentCell
+    ElmanCell
+    LSTMCell
+    GRUCell
+
+Attention and transformers
+--------------------------
+
+.. autosummary::
+    :toctree: generated/
+
+    MultiheadAttention
+    CausalSelfAttention
+    FeedForward
+    TransformerBlock
+    scaled_dot_product_attention

pytensor_ml-0.3.0/docs/source/api/loss.rst

@@ -0,0 +1,12 @@
+Losses
+======
+
+.. currentmodule:: pytensor_ml.loss
+
+.. autosummary::
+    :toctree: generated/
+
+    Loss
+    SquaredError
+    CrossEntropy
+    supervised_loss

pytensor_ml-0.3.0/docs/source/api/model.rst

@@ -0,0 +1,9 @@
+Model
+=====
+
+.. currentmodule:: pytensor_ml.model
+
+.. autosummary::
+    :toctree: generated/
+
+    Model