pytensor-ml 0.0.4__tar.gz → 0.0.6__tar.gz

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

@@ -0,0 +1,57 @@
+name: mypy
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+    paths:
+      - ".github/workflows/mypy.yml"
+      - "pytensor_ml/**"
+      - "scripts/run_mypy.py"
+      - "pyproject.toml"
+      - "conda_envs/**"
+
+# 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
+
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+    - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+    - uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+      env:
+        CACHE_NUMBER: 0
+      with:
+        path: ~/conda_pkgs_dir
+        key: ${{ runner.os }}-py312-conda-${{ env.CACHE_NUMBER }}-${{
+          hashFiles('conda_envs/pytensor_ml.yml') }}
+    - 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
+        use-only-tar-bz2: false
+
+    - name: Install current branch
+      run: |
+        conda activate pytensor_ml
+        pip install -e .
+        python --version
+
+    - name: Run mypy
+      run: |
+        conda activate pytensor_ml
+        python ./scripts/run_mypy.py

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/.github/workflows/release.yml

@@ -30,7 +30,7 @@ jobs:
 
   publish-package:
     # Don't publish from forks
-    if: github.repository_owner == 'jessegrabowski' && github.event_name == 'release' && github.event.action == 'published'
+    if: github.repository_owner == 'pymc-devs' && github.event_name == 'release' && github.event.action == 'published'
     # Use the `release` GitHub environment to protect the Trusted Publishing (OIDC)
     # workflow by requiring signoff from a maintainer.
     environment: release

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/.pre-commit-config.yaml

@@ -12,7 +12,7 @@ repos:
     -   id: trailing-whitespace
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.15.15
+  rev: v0.15.17
   hooks:
     - id: ruff
       args: [ --fix, --exit-non-zero-on-fix ]
@@ -50,7 +50,7 @@ repos:
                conf.py)
       entry: >
           (?x)(arviz-devs.github.io|
-               pymc-experimental.readthedocs.io|
+               pymc-extras.readthedocs.io|
                docs.pymc.io|
                numpy.org/doc|
                pymc-examples.readthedocs.io|

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: pytensor_ml
-Version: 0.0.4
+Version: 0.0.6
 Summary: A(nother) deep learning library, built on top of Pytensor
-Project-URL: Homepage, https://github.com/jessegrabowski/pytensor_ml
-Project-URL: Source, https://github.com/jessegrabowski/pytensor_ml
-Project-URL: Tracker, https://github.com/jessegrabowski/pytensor_ml/issues
+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:                                  Apache License
@@ -224,9 +224,11 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering
 Requires-Python: >=3.12
 Requires-Dist: numpy
-Requires-Dist: pymc<7.0.0,>=6.0.0
-Requires-Dist: pytensor<4.0.0,>=3.0.0
+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'

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/conda_envs/pytensor_ml-gpu_jax.yml

@@ -6,8 +6,7 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.0.0,<4.0.0
-  - pymc>=6.0.0,<7.0.0
+  - pytensor>=3.2.3,<4.0.0
   - numpy
   - scikit-learn
 

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/conda_envs/pytensor_ml.yml

@@ -5,9 +5,9 @@ channels:
 
 dependencies:
   - python>=3.12
-  - pytensor>=3.0.0,<4.0.0
-  - pymc>=6.0.0,<7.0.0
+  - pytensor>=3.2.3,<4.0.0
   - numpy
+  - safetensors
   - scikit-learn
 
   # For testing
@@ -28,3 +28,9 @@ dependencies:
 
   # developer tools
   - pre-commit
+  - mypy
+
+  # JAX backend for dispatch tests; via pip so Windows pulls CPU wheels instead of conda jaxlib
+  - pip
+  - pip:
+      - jax

pytensor_ml-0.0.6/examples/mnist_feed_forward.ipynb

@@ -0,0 +1,198 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15279e7f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00915a61",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pytensor\n",
+    "\n",
+    "floatX = \"float32\"\n",
+    "pytensor.config.floatX = floatX"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fae24c05",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "from sklearn.datasets import load_digits\n",
+    "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.loss import CrossEntropy\n",
+    "from pytensor_ml.model import Model\n",
+    "from pytensor_ml.optim import adagrad\n",
+    "from pytensor_ml.util import DataLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "65882b64",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "X, y = load_digits(return_X_y=True)\n",
+    "y_onehot = OneHotEncoder().fit_transform(y[:, None]).toarray().astype(floatX)\n",
+    "X_normed = MinMaxScaler().fit_transform(X).astype(floatX)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f81f9076",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# X_in = Input('X_in', shape=(64, ))\n",
+    "X_in = pytensor.tensor.tensor(\"X_in\", shape=(None, 64))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "41840fac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prediction_network = Sequential(\n",
+    "    Linear(\"Linear_1\", n_in=64, n_out=256),\n",
+    "    BatchNorm2D(\"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",
+    "    Dropout(\"Dropout_2\", p=0.5),\n",
+    "    LeakyReLU(),\n",
+    "    Linear(\"Logits\", n_in=128, n_out=10),\n",
+    ")\n",
+    "\n",
+    "y_hat = prediction_network(X_in)\n",
+    "model = Model(X_in, y_hat, compile_kwargs={\"mode\": \"NUMBA\"})\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3bfa0ddc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loss_fn = CrossEntropy(expect_onehot_labels=True, expect_logits=True, reduction=\"mean\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "290d3191",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# `model.compile_train` is the training companion to `model.predict`: it builds the loss against a target,\n",
+    "# differentiates it, folds in the BatchNorm running-statistic updates, and compiles a one-step training\n",
+    "# function (reusing the model's NUMBA compile settings).\n",
+    "step = model.compile_train(adagrad(learning_rate=1e-3), loss_fn, ndim_out=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e77b697f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.initialize()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd6dede5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tqdm.notebook import tqdm\n",
+    "\n",
+    "loader = DataLoader(X_normed, y_onehot)\n",
+    "loss_history = []\n",
+    "n_epochs = 10_000\n",
+    "\n",
+    "for _ in tqdm(range(n_epochs)):\n",
+    "    X_batch, y_batch = loader()\n",
+    "    loss_value = step(X_batch, y_batch)\n",
+    "    loss_history.append(loss_value)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab1ad6ef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.plot(loss_history)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb38dd99",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from scipy.special import softmax\n",
+    "\n",
+    "# `model.predict` compiles a forward pass specialized for inference: Dropout is dropped and BatchNorm uses its\n",
+    "# running statistics.\n",
+    "y_hat_logits = model.predict(X_normed)\n",
+    "y_hat_probs = softmax(y_hat_logits, axis=-1)\n",
+    "y_pred = np.argmax(y_hat_probs, axis=-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4393b9c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import seaborn as sns\n",
+    "\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "sns.heatmap(confusion_matrix(y, y_pred), annot=True, fmt=\"0.0f\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/pyproject.toml

@@ -34,22 +34,24 @@ keywords = [
 ]
 
 dependencies = [
-  "pytensor>=3.0.0,<4.0.0",
+  "pytensor>=3.2.3,<4.0.0",
   "numpy",
-  "pymc>=6.0.0,<7.0.0"
+  "safetensors"
 ]
 
 [project.urls]
-Homepage = "https://github.com/jessegrabowski/pytensor_ml"
-Source = "https://github.com/jessegrabowski/pytensor_ml"
-Tracker = "https://github.com/jessegrabowski/pytensor_ml/issues"
+Homepage = "https://github.com/pymc-devs/pytensor_ml"
+Source = "https://github.com/pymc-devs/pytensor_ml"
+Tracker = "https://github.com/pymc-devs/pytensor_ml/issues"
 
 [project.optional-dependencies]
 dev = [
     "pre-commit",
     "pytest",
     "pytest-cov",
-    "scikit-learn"
+    "scikit-learn",
+    "mypy",
+    "jax",
 ]
 
 [tool.hatch.version]
@@ -109,6 +111,24 @@ ignore = [
 ]
 fixable = ["ALL"]
 
+[tool.ruff.lint.per-file-ignores]
+# An __init__.py re-exports its package's API and imports submodules for their registration side effects.
+"**/__init__.py" = ["F401"]
+
 
 [tool.ruff.lint.isort]
 lines-between-types = 1
+
+
+[tool.mypy]
+python_version = "3.12"
+ignore_missing_imports = true
+strict_equality = true
+warn_redundant_casts = true
+warn_unused_configs = true
+warn_unused_ignores = true
+# pytensor ops are typed to return Any, so returning their results from a function annotated with a
+# concrete tensor type fires constantly without catching real bugs in our code.
+warn_return_any = false
+warn_unreachable = true
+files = ["pytensor_ml"]

pytensor_ml-0.0.6/pytensor_ml/__init__.py

@@ -0,0 +1,29 @@
+import logging
+import sys
+
+import pytensor_ml.dispatch  # installs the lazy backend-dispatch registration hook
+
+from pytensor_ml._version import __version__
+from pytensor_ml.checkpoint import load_state, save_state
+from pytensor_ml.pretrained import from_pretrained, load_network, save_network, save_pretrained
+from pytensor_ml.pytensorf import function
+
+_log = logging.getLogger(__name__)
+
+if not logging.root.handlers:
+    _log.setLevel(logging.INFO)
+    if len(_log.handlers) == 0:
+        handler = logging.StreamHandler(sys.stderr)
+        _log.addHandler(handler)
+
+
+__all__ = [
+    "__version__",
+    "from_pretrained",
+    "function",
+    "load_network",
+    "load_state",
+    "save_network",
+    "save_pretrained",
+    "save_state",
+]

{pytensor_ml-0.0.4 → pytensor_ml-0.0.6}/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.4'
-__version_tuple__ = version_tuple = (0, 0, 4)
+__version__ = version = '0.0.6'
+__version_tuple__ = version_tuple = (0, 0, 6)
 
 __commit_id__ = commit_id = None

pytensor_ml-0.0.6/pytensor_ml/activations.py

@@ -0,0 +1,225 @@
+import numpy as np
+import pytensor.tensor as pt
+
+from pytensor import config
+
+from pytensor_ml.base import Layer
+
+
+def _constant_like(value: float, x: pt.TensorVariable) -> pt.TensorVariable:
+    """
+    Wrap a scalar so that combining it with ``x`` cannot widen ``x``'s dtype.
+
+    PyTensor's autocaster types a bare Python float by value, so whether a literal widens its operand
+    depends on that value: against a float32 input ``0.5 * x`` stays float32, while ``0.01 * x``
+    promotes to float64. Pinning the constant to ``x``'s dtype removes the dependence.
+    """
+    dtype = np.dtype(x.dtype)
+    # np.finfo maps complex64 -> float32, keeping complex inputs at their own precision.
+    dtype = np.finfo(dtype).dtype if np.issubdtype(dtype, np.inexact) else np.dtype(config.floatX)
+    return pt.constant(np.asarray(value, dtype=dtype))
+
+
+class Activation(Layer): ...
+
+
+class ReLU(Activation):
+    r"""
+    Rectified Linear Unit.
+
+    Compute the positive part of the input:
+
+    .. math::
+
+        \mathrm{ReLU}(x) = \max(0, x).
+    """
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.maximum(0, x)
+        out.name = "ReLU"
+        return out
+
+
+class LeakyReLU(Activation):
+    r"""
+    Leaky Rectified Linear Unit.
+
+    Replace ReLU's flat negative branch with a small negative slope, so negative inputs keep a
+    nonzero gradient:
+
+    .. math::
+
+        \mathrm{LeakyReLU}(x) = \begin{cases} x & x > 0 \\ \alpha x & x \le 0 \end{cases}
+
+    Parameters
+    ----------
+    negative_slope : float, optional
+        The slope :math:`\alpha` applied to negative inputs. Default is 0.01.
+    """
+
+    def __init__(self, negative_slope: float = 0.01):
+        self.negative_slope = negative_slope
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        x = pt.as_tensor(x)
+        out = pt.switch(x > 0, x, _constant_like(self.negative_slope, x) * x)
+        out.name = "LeakyReLU"
+        return out
+
+
+class Tanh(Activation):
+    r"""
+    Hyperbolic tangent.
+
+    Squash the input to :math:`(-1, 1)`:
+
+    .. math::
+
+        \tanh(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}.
+    """
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.tanh(x)
+        out.name = "TanH"
+        return out
+
+
+class Sigmoid(Activation):
+    r"""
+    Logistic sigmoid.
+
+    Squash the input to :math:`(0, 1)`:
+
+    .. math::
+
+        \sigma(x) = \frac{1}{1 + e^{-x}}.
+    """
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.sigmoid(x)
+        out.name = "Sigmoid"
+        return out
+
+
+class SoftPlus(Activation):
+    r"""
+    Softplus activation.
+
+    Compute a smooth approximation to ReLU:
+
+    .. math::
+
+        \mathrm{softplus}(x) = \log(1 + e^x).
+    """
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.softplus(x)
+        out.name = "SoftPlus"
+        return out
+
+
+class GELU(Activation):
+    r"""
+    Gaussian Error Linear Unit.
+
+    Compute :math:`\mathrm{GELU}(x) = x \, \Phi(x)`, where :math:`\Phi` is the standard normal
+    cumulative distribution function:
+
+    .. math::
+
+        \mathrm{GELU}(x) = \frac{x}{2} \left(1 + \operatorname{erf}\!\left(\frac{x}{\sqrt{2}}\right)\right).
+
+    Parameters
+    ----------
+    approximate : bool, optional
+        Use the tanh approximation
+
+        .. math::
+
+            \mathrm{GELU}(x) \approx \frac{x}{2}
+            \left(1 + \tanh\!\left[\sqrt{2/\pi}\,(x + 0.044715\,x^3)\right]\right)
+
+        This is the variant HuggingFace calls ``"gelu_new"`` / ``"gelu_pytorch_tanh"``, PyTorch exposes
+        as ``nn.GELU(approximate="tanh")``, and Flax as ``gelu(approximate=True)``; GPT-2 uses it. It is
+        cheaper to evaluate than the exact :math:`\operatorname{erf}` form. Default is True.
+    """
+
+    def __init__(self, approximate: bool = True):
+        self.approximate = approximate
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        x = pt.as_tensor(x)
+        half = _constant_like(0.5, x)
+        one = _constant_like(1.0, x)
+        if self.approximate:
+            sqrt_2_over_pi = _constant_like(np.sqrt(2.0 / np.pi), x)
+            cubic_coef = _constant_like(0.044715, x)
+            out = half * x * (one + pt.tanh(sqrt_2_over_pi * (x + cubic_coef * x**3)))
+        else:
+            sqrt2 = _constant_like(np.sqrt(2.0), x)
+            out = half * x * (one + pt.erf(x / sqrt2))
+        out.name = "GELU"
+        return out
+
+
+class Swish(Activation):
+    r"""
+    Swish activation, also known as SiLU.
+
+    Compute :math:`\mathrm{Swish}(x) = x \, \sigma(\beta x)`, where :math:`\sigma` is the logistic
+    sigmoid. With :math:`\beta = 1` this is the Sigmoid Linear Unit (PyTorch ``nn.SiLU``, HuggingFace
+    ``"silu"`` / ``"swish"``).
+
+    Parameters
+    ----------
+    beta : float, optional
+        Slope of the sigmoid gate. Larger :math:`\beta` sharpens the gate toward a ReLU; :math:`\beta
+        \to 0` collapses it toward the linear map :math:`x/2`. Default is 1.0.
+    """
+
+    def __init__(self, beta: float = 1.0):
+        self.beta = beta
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        x = pt.as_tensor(x)
+        out = x * pt.sigmoid(_constant_like(self.beta, x) * x)
+        out.name = "Swish"
+        return out
+
+
+class Softmax(Activation):
+    r"""
+    Softmax activation.
+
+    Normalize the input along one axis into a probability distribution:
+
+    .. math::
+
+        \mathrm{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}.
+
+    Parameters
+    ----------
+    axis : int, optional
+        The axis along which the values sum to one. Default is -1.
+    """
+
+    def __init__(self, axis: int = -1):
+        self.axis = axis
+
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.special.softmax(x, axis=self.axis)
+        out.name = "Softmax"
+        return out
+
+
+__all__ = [
+    "GELU",
+    "Activation",
+    "LeakyReLU",
+    "ReLU",
+    "Sigmoid",
+    "SoftPlus",
+    "Softmax",
+    "Swish",
+    "Tanh",
+]