pytensor-ml 0.0.1__tar.gz → 0.0.2__tar.gz

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor_ml
-Version: 0.0.1
+Version: 0.0.2
 Summary: A(nother) deep learning library, built on top of Pytensor
 Author: Ricardo Vieira
 Author-email: Jesse Grabowski <jessegrabowski@gmail.com>

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pyproject.toml

@@ -33,7 +33,8 @@ keywords = [
 
 dependencies = [
   "pytensor",
-  "numpy"
+  "numpy",
+  "pymc"
 ]
 
 [project.optional-dependencies]
@@ -96,6 +97,8 @@ ignore = [
   "D415",
   "D417",
 ]
+fixable = ["ALL"]
+
 
 [tool.ruff.lint.isort]
 lines-between-types = 1

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml/__init__.py

@@ -2,6 +2,7 @@ import logging
 import sys
 
 from pytensor_ml._version import get_versions
+from pytensor_ml.pytensorf import function
 
 _log = logging.getLogger(__name__)
 
@@ -13,3 +14,5 @@ if not logging.root.handlers:
 
 
 __version__ = get_versions()["version"]
+
+__all__ = ["function"]

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-30T10:40:07+0800",
+ "date": "2025-01-26T00:24:12+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "599c7aacb627660af38618e6095d8762a8d49a0a",
- "version": "0.0.1"
+ "full-revisionid": "75eda5beff184dd32a4fe21444ec961a27c8eca3",
+ "version": "0.0.2"
 }
 '''  # END VERSION_JSON
 

pytensor_ml-0.0.2/pytensor_ml/layers.py

@@ -0,0 +1,103 @@
+from abc import ABC
+from collections.abc import Callable
+
+import numpy as np
+import pytensor.tensor as pt
+import pytensor.tensor.random as ptr
+
+from pytensor import config
+from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.sharedvalue import shared
+
+
+def shape_to_str(shape):
+    inner = ",".join([str(st_dim) if st_dim is not None else "?" for st_dim in shape])
+    return f"({inner})"
+
+
+class Layer(ABC):
+    def __call__(self, x: pt.TensorLike) -> pt.TensorLike: ...
+
+
+class LinearLayer(OpFromGraph): ...
+
+
+class Linear(Layer):
+    __props__ = ("name", "n_in", "n_out")
+
+    def __init__(self, name: str | None, n_in: int, n_out: int):
+        self.name = name if name else "Linear"
+        self.n_in = n_in
+        self.n_out = n_out
+
+        self.W = pt.tensor(f"{self.name}_W", shape=(n_in, self.n_out))
+        self.b = pt.tensor(f"{self.name}_b", shape=(self.n_out,))
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorLike:
+        X = pt.as_tensor(X)
+
+        init_st_shape = shape_to_str(X.type.shape)
+
+        res = X @ self.W + self.b
+
+        final_st_shape = shape_to_str(res.type.shape)
+
+        ofg = LinearLayer(
+            inputs=[X, self.W, self.b],
+            outputs=[res],
+            inline=True,
+            name=f"{self.name}[{init_st_shape} -> {final_st_shape}]",
+        )
+        out = ofg(X, self.W, self.b)
+        out.name = f"{self.name}_output"
+
+        return out
+
+
+class DropoutLayer(OpFromGraph): ...
+
+
+class Dropout(Layer):
+    __props__ = ("name", "p")
+
+    def __init__(self, name: str | None, p: float = 0.5):
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"Dropout probability has to be between 0 and 1, but got {p}")
+        self.name = name if name else "Dropout"
+        self.p = p
+        self.rng = shared(np.random.default_rng())
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorLike:
+        X = pt.as_tensor(X)
+        new_rng, mask = ptr.bernoulli(p=1 - self.p, size=X.shape, rng=self.rng).owner.outputs
+        mask = mask.astype(config.floatX)
+
+        X_masked = DropoutLayer(
+            inputs=[X, mask], outputs=[pt.where(mask, ift=X / (1 - self.p), iff=0)], inline=True
+        )(X, mask)
+        X_masked.name = f"{self.name}[p = {self.p}]"
+
+        return X_masked
+
+
+def Input(name: str, shape: tuple[int]) -> pt.TensorLike:
+    if not all(isinstance(dim, int) for dim in shape):
+        raise ValueError("All dimensions must be integers")
+
+    return pt.tensor(name=name, shape=shape)
+
+
+def Sequential(*layers: Callable) -> Callable:
+    def forward(x: pt.TensorLike) -> pt.TensorLike:
+        for layer in layers:
+            x = layer(x)
+        return x
+
+    return forward
+
+
+Squeeze = pt.squeeze
+Concatenate = pt.concatenate
+
+
+__all__ = ["Concatenate", "Input", "Linear", "Sequential", "Squeeze"]

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml/model.py

@@ -5,13 +5,14 @@ from typing import Literal, cast
 import numpy as np
 
 from pytensor import config
-from pytensor.compile.function import function
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph import graph_inputs
 from pytensor.graph.basic import Constant
 from pytensor.printing import debugprint
 from pytensor.tensor import TensorVariable
 
+from pytensor_ml.pytensorf import function, rewrite_for_prediction
+
 InitializationSchemes = Literal["zeros", "xavier_uniform", "xavier_normal"]
 
 
@@ -24,14 +25,15 @@ def required_graph_inputs(tensor: TensorVariable) -> Generator[TensorVariable, N
 
 
 class Model:
-    def __init__(self, X: TensorVariable, y: TensorVariable):
+    def __init__(self, X: TensorVariable, y: TensorVariable, compile_kwargs: dict | None = None):
         self.X = X
         self.y = y
 
+        self._compile_kwargs = compile_kwargs if compile_kwargs else {}
         self._weight_values: list[np.ndarray[float]] | None = None
         self._predict_fn = None
 
-    def initalize_weights(
+    def initialize_weights(
         self,
         scheme: InitializationSchemes = "xavier_normal",
         random_seed: int | str | np.random.Generator | None = None,
@@ -55,10 +57,14 @@ class Model:
 
     def predict(self, X_values: np.ndarray) -> np.ndarray:
         if self._predict_fn is None:
-            f = function([*self.weights, self.X], self.y)
+            f = function(
+                [*self.weights, self.X],
+                rewrite_for_prediction(self.y),
+                **self._compile_kwargs,
+            )
             self._predict_fn = partial(f, *self.weight_values)
 
-        return self._predict_fn(X_values)
+        return cast(np.ndarray, self._predict_fn(X_values))
 
     def __str__(self):
         return debugprint(self.y, file="str")

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml/optimizers.py

@@ -4,12 +4,12 @@ from collections.abc import Callable, Sequence
 import numpy as np
 
 from pytensor import config
-from pytensor.compile.function import function
 from pytensor.compile.function.types import Function
 from pytensor.gradient import grad
 from pytensor.tensor import TensorLike, TensorVariable, sqrt, tensor
 
 from pytensor_ml.model import Model
+from pytensor_ml.pytensorf import function, rewrite_pregrad
 
 
 class Optimizer(ABC):
@@ -19,6 +19,7 @@ class Optimizer(ABC):
         loss_fn: Callable,
         ndim_out: int = 1,
         optimizer_weights: list[TensorVariable] | None = None,
+        compile_kwargs: dict | None = None,
     ):
         self.model = model
         self.loss_fn = loss_fn
@@ -26,7 +27,7 @@ class Optimizer(ABC):
         self.optimizer_weights = optimizer_weights if optimizer_weights is not None else []
         self._optimizer_weights_values = self._initialize_weights()
 
-        self.update_fn = self.build_update_fn()
+        self.update_fn = self.build_update_fn(compile_kwargs)
 
     @abstractmethod
     def update_parameters(self, params: Sequence[TensorVariable], loss: TensorVariable): ...
@@ -51,7 +52,7 @@ class Optimizer(ABC):
         n_params = len(self.model.weights)
         return all_weights[:n_params], all_weights[n_params:]
 
-    def build_update_fn(self) -> Function:
+    def build_update_fn(self, compile_kwargs) -> Function:
         """
         Compile a function to update model weights
 
@@ -59,12 +60,17 @@ class Optimizer(ABC):
         values. By default, the functions also returns the loss value. To return additional diagnostics, this method
         should be overloaded.
 
+        Parameters
+        ----------
+        compile_kwargs: dict
+            Keyword arguments passed to pytensor.function
+
         Returns
         -------
         update_fn: Function
             A function that updates the model weights given a batch of data, targets, and current weights.
         """
-
+        compile_kwargs = compile_kwargs if compile_kwargs else {}
         x, y_hat = self.model.X, self.model.y
 
         label_slice = (slice(None),) * self.ndim_out + (0,) * (y_hat.ndim - self.ndim_out)
@@ -74,9 +80,16 @@ class Optimizer(ABC):
         optimizer_weights = self.optimizer_weights
 
         loss = self.loss_fn(y, y_hat)
+        loss = rewrite_pregrad(loss)
 
         all_weights = self.update_parameters(weights + optimizer_weights, loss)
-        fn = function([x, y, *weights, *optimizer_weights], [*all_weights, loss])
+        fn = function(
+            [x, y, *weights, *optimizer_weights],
+            [*all_weights, loss],
+            **compile_kwargs,
+        )
+        fn.trust_input = True
+
         return fn
 
     def step(self, x_values, y_values) -> np.ndarray:
@@ -90,7 +103,6 @@ class Optimizer(ABC):
         *new_weights, loss_values = self.update_fn(
             x_values, y_values, *self.model.weight_values, *self.optimizer_weights_values
         )
-        self.update_fn.trust_input = True
 
         self.model.weight_values, self.optimizer_weights_values = self._split_weights(new_weights)
 
@@ -98,9 +110,17 @@ class Optimizer(ABC):
 
 
 class SGD(Optimizer):
-    def __init__(self, model, loss_fn, *, ndim_out: int = 1, learning_rate: TensorLike = 0.01):
+    def __init__(
+        self,
+        model,
+        loss_fn,
+        *,
+        ndim_out: int = 1,
+        learning_rate: TensorLike = 0.01,
+        compile_kwargs: dict | None = None,
+    ):
         self.learning_rate = learning_rate
-        super().__init__(model, loss_fn, ndim_out=ndim_out)
+        super().__init__(model, loss_fn, ndim_out=ndim_out, compile_kwargs=compile_kwargs)
 
     def update_parameters(
         self, params: list[TensorVariable], loss: TensorVariable
@@ -122,17 +142,25 @@ class ADAGrad(Optimizer):
         ndim_out: int = 1,
         learning_rate: TensorLike = 0.01,
         epsilon: TensorLike = 1e-8,
+        compile_kwargs: dict | None = None,
     ):
         self.learning_rate = learning_rate
         self.epsilon = epsilon
         g2_weights = [param.type() for param in model.weights]
 
-        super().__init__(model, loss_fn, ndim_out=ndim_out, optimizer_weights=g2_weights)
+        super().__init__(
+            model,
+            loss_fn,
+            ndim_out=ndim_out,
+            optimizer_weights=g2_weights,
+            compile_kwargs=compile_kwargs,
+        )
 
     def update_parameters(
         self, weights: list[TensorVariable], loss: TensorVariable
     ) -> list[TensorVariable]:
         weights, optimizer_weights = self._split_weights(weights)
+
         grads = grad(loss, weights)
 
         new_weights = []
@@ -158,6 +186,7 @@ class Adam(Optimizer):
         beta1: TensorLike = 0.9,
         beta2: TensorLike = 0.999,
         epsilon: TensorLike = 1e-8,
+        compile_kwargs: dict | None = None,
     ):
         self.learning_rate = learning_rate
         self.beta1 = beta1
@@ -169,7 +198,13 @@ class Adam(Optimizer):
         t = tensor("t", shape=(1,), dtype=config.floatX)
 
         optimizer_weights = m_weights + v_weights + [t]
-        super().__init__(model, loss_fn, ndim_out=ndim_out, optimizer_weights=optimizer_weights)
+        super().__init__(
+            model,
+            loss_fn,
+            ndim_out=ndim_out,
+            optimizer_weights=optimizer_weights,
+            compile_kwargs=compile_kwargs,
+        )
 
     def update_parameters(
         self, weights: list[TensorVariable], loss: TensorVariable

pytensor_ml-0.0.2/pytensor_ml/pytensorf.py

@@ -0,0 +1,95 @@
+from typing import cast
+
+import pytensor
+
+from pymc.pytensorf import SeedSequenceSeed, collect_default_updates, reseed_rngs
+from pytensor import Mode
+from pytensor.compile import Function, SharedVariable, get_mode
+from pytensor.graph import FunctionGraph, RewriteDatabaseQuery, rewrite_graph
+from pytensor.tensor.variable import Variable
+
+from pytensor_ml.rewriting.layers import predict_db
+
+
+def atleast_list(x):
+    if not isinstance(x, list | tuple):
+        return [x]
+    return x
+
+
+def function(
+    inputs,
+    outputs,
+    random_seed: SeedSequenceSeed = None,
+    mode=None,
+    **kwargs,
+) -> Function:
+    """
+    Compile a Pytensor function, including specialized rewrites.
+
+    Parameters
+    ----------
+    inputs: list of TensorVariables, optional
+        Inputs of the compiled PyTensor function
+    outputs: list of TensorVariables, optional
+        Outputs of the compiled PyTensor function
+    random_seed: int, array-like of int or SeedSequence, optional
+        Seed used to override any RandomState/Generator shared variables in the graph.
+        If not specified, the value of original shared variables will still be overwritten.
+    mode: optional
+        PyTensor mode used to compile the function
+
+    Returns
+    -------
+    pytensor_function: Function
+        Compiled function
+    """
+    rng_updates = collect_default_updates(
+        inputs=[inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs],
+        outputs=[
+            out.variable if isinstance(out, pytensor.Out) else out for out in atleast_list(outputs)
+        ],
+    )
+
+    if rng_updates:
+        rngs = cast(list[SharedVariable], list(rng_updates))
+        reseed_rngs(rngs, random_seed)
+
+    mode = get_mode(mode)
+    opt_qry = mode.provided_optimizer.including("random_make_inplace")
+    mode = Mode(linker=mode.linker, optimizer=opt_qry)
+    pytensor_function = pytensor.function(
+        inputs,
+        outputs,
+        updates={**rng_updates, **kwargs.pop("updates", {})},
+        mode=mode,
+        **kwargs,
+    )
+    return pytensor_function
+
+
+def rewrite_pregrad(graph):
+    """Apply simplifying or stabilizing rewrites to graph that are safe to use pre-grad."""
+    return rewrite_graph(graph, include=("canonicalize", "stabilize"))
+
+
+def rewrite_for_prediction(graph):
+    """Apply rewrites to specialize a graph for forward passes (e.g. removing Dropout layers)"""
+    if isinstance(graph, FunctionGraph):
+        fgraph = graph
+    else:
+        outputs = [graph] if isinstance(graph, Variable) else graph
+        fgraph = FunctionGraph(outputs=outputs, clone=False)
+
+    rewriter = predict_db.query(RewriteDatabaseQuery(include=["basic"]))
+    rewriter.rewrite(fgraph)
+
+    if isinstance(graph, FunctionGraph):
+        return fgraph
+    if isinstance(graph, Variable):
+        return fgraph.outputs[0]
+
+    return fgraph.outputs
+
+
+__all__ = ["function"]

pytensor_ml-0.0.2/pytensor_ml/rewriting/layers.py

@@ -0,0 +1,37 @@
+from pytensor.graph.basic import Apply
+from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.rewriting.basic import node_rewriter
+from pytensor.graph.rewriting.db import EquilibriumDB
+from pytensor.tensor.variable import Variable
+
+from pytensor_ml.layers import DropoutLayer
+
+predict_db = EquilibriumDB()
+
+
+@node_rewriter([DropoutLayer])
+def remove_dropout_for_prediction(fgraph: FunctionGraph, node: Apply) -> list[Variable] | None:
+    """
+    Set dropout probability to zero for all dropout layers.
+
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        Graph being rewritten
+    node: Node
+        Node being rewritten
+
+    Returns
+    -------
+    X: Variable
+        The input to the dropout layer, removing the dropout from the graph
+    """
+    X, rng = node.inputs
+    return [X]
+
+
+predict_db.register(
+    "remove_dropout_for_prediction",
+    remove_dropout_for_prediction,
+    "basic",
+)

pytensor_ml-0.0.2/pytensor_ml/util.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+from pytensor import config
+
+
+class DataLoader:
+    def __init__(self, X, y, batch_size=64, dtype=None):
+        if dtype is None:
+            dtype = config.floatX
+        self.rng = np.random.default_rng()
+
+        self.X = X.astype(dtype)
+        self.y = y.astype(dtype)
+        self.n = X.shape[0]
+
+        self.batch_size = batch_size
+        self.cursor = 0
+        self.indices = np.arange(len(X)).astype("int32")
+        self.reset()
+
+    def shuffle(self):
+        self.rng.shuffle(self.indices)
+
+    def move_cursor(self):
+        start, stop = self.cursor, self.cursor + self.batch_size
+
+        has_remainder, stop = divmod(stop, self.n)
+
+        idx = slice(start, None if has_remainder > 0 else stop)
+        indices = self.indices[idx]
+
+        if has_remainder > 0:
+            excess = self.batch_size - indices.shape[0]
+            self.shuffle()
+            self.epoch += 1
+            indices = np.r_[indices, self.indices[:excess]]
+
+        self.cursor = (self.cursor + self.batch_size) % self.n
+
+        return indices
+
+    def reset(self):
+        self.cursor = 0
+        self.epoch = 0
+        self.shuffle()
+
+    def __call__(self):
+        idx = self.move_cursor()
+        return self.X[idx], self.y[idx]

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor-ml
-Version: 0.0.1
+Version: 0.0.2
 Summary: A(nother) deep learning library, built on top of Pytensor
 Author: Ricardo Vieira
 Author-email: Jesse Grabowski <jessegrabowski@gmail.com>

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml.egg-info/SOURCES.txt

@@ -9,8 +9,12 @@ pytensor_ml/layers.py
 pytensor_ml/loss.py
 pytensor_ml/model.py
 pytensor_ml/optimizers.py
+pytensor_ml/pytensorf.py
+pytensor_ml/util.py
 pytensor_ml.egg-info/PKG-INFO
 pytensor_ml.egg-info/SOURCES.txt
 pytensor_ml.egg-info/dependency_links.txt
 pytensor_ml.egg-info/requires.txt
-pytensor_ml.egg-info/top_level.txt
+pytensor_ml.egg-info/top_level.txt
+pytensor_ml/rewriting/__init__.py
+pytensor_ml/rewriting/layers.py

{pytensor_ml-0.0.1 → pytensor_ml-0.0.2}/pytensor_ml.egg-info/requires.txt

@@ -1,5 +1,6 @@
 pytensor
 numpy
+pymc
 
 [dev]
 pre-commit

pytensor_ml-0.0.1/pytensor_ml/layers.py

@@ -1,53 +0,0 @@
-from abc import ABC
-from collections.abc import Callable
-
-from pytensor.tensor import TensorLike, as_tensor, concatenate, squeeze, tensor
-
-
-class Layer(ABC):
-    def __call__(self, x: TensorLike) -> TensorLike: ...
-
-
-class Linear(Layer):
-    __props__ = ("name", "n_out")
-
-    def __init__(self, name: str | None, n_in: int, n_out: int):
-        self.name = name
-        self.n_in = n_in
-        self.n_out = n_out
-
-        self.W = tensor(f"{self.name}_W", shape=(n_in, self.n_out))
-        self.b = tensor(f"{self.name}_b", shape=(self.n_out,))
-
-    def __call__(self, X: TensorLike) -> TensorLike:
-        X = as_tensor(X)
-
-        init_st_shape = [st_dim if st_dim is not None else "?" for st_dim in X.type.shape]
-        X = X @ self.W + self.b
-        final_st_shape = [st_dim if st_dim is not None else "?" for st_dim in X.type.shape]
-        X.name = f"{self.__class__.__name__}[{init_st_shape} -> {final_st_shape}]"
-
-        return X
-
-
-def Input(name: str, shape: tuple[int]) -> TensorLike:
-    if not all(isinstance(dim, int) for dim in shape):
-        raise ValueError("All dimensions must be integers")
-
-    return tensor(name=name, shape=shape)
-
-
-def Sequential(*layers: Callable) -> Callable:
-    def forward(x: TensorLike) -> TensorLike:
-        for layer in layers:
-            x = layer(x)
-        return x
-
-    return forward
-
-
-Squeeze = squeeze
-Concatenate = concatenate
-
-
-__all__ = ["Concatenate", "Input", "Linear", "Sequential", "Squeeze"]