Add a Sphinx docs site and an example on every public entrypoint

.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"

.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

README.md

@@ -1,2 +1,96 @@
 # pytensor_ml
-Neural network package built on Pytensor
+
+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`, `numpy`, and `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(X_in, network(X_in)).initialize(seed=0)
+
+rule = chain(adam(learning_rate=cosine_schedule(1e-3, total_steps=500)), clip_by_global_norm(1.0))
+loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction="mean")
+step = model.compile_train(rule, loss_fn, ndim_out=2)
+
+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).

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.0,<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 ..

conda_envs/pytensor_ml.yml

@@ -20,13 +20,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

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/

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)

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

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 %}

docs/source/api.rst

@@ -0,0 +1,18 @@
+.. _api:
+
+API Reference
+=============
+
+.. toctree::
+    :maxdepth: 1
+    :titlesonly:
+
+    api/model
+    api/layers
+    api/activations
+    api/loss
+    api/optim
+    api/state
+    api/serialization
+    api/pytensorf
+    api/util

docs/source/api/activations.rst

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