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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor_ml
-Version: 0.0.2
+Version: 0.0.3
 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.2 → pytensor_ml-0.0.3}/pytensor_ml/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-01-26T00:24:12+0800",
+ "date": "2025-06-18T22:13:49+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "75eda5beff184dd32a4fe21444ec961a27c8eca3",
- "version": "0.0.2"
+ "full-revisionid": "56994236bd76d029096eb5ad6ca3767a893bb4fa",
+ "version": "0.0.3"
 }
 '''  # END VERSION_JSON
 

{pytensor_ml-0.0.2 → pytensor_ml-0.0.3}/pytensor_ml/activations.py

@@ -37,6 +37,13 @@ class Sigmoid(Activation):
         return out
 
 
+class SoftPlus(Activation):
+    def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
+        out = pt.softplus(x)
+        out.name = "SoftPlus"
+        return out
+
+
 class Softmax(Activation):
     def __init__(self, axis: int = -1):
         self.axis = axis

pytensor_ml-0.0.3/pytensor_ml/layers.py

@@ -0,0 +1,236 @@
+from abc import ABC
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import pytensor.tensor as pt
+import pytensor.tensor.random as ptr
+
+from pytensor import config
+from pytensor.compile.sharedvalue import shared
+
+from pytensor_ml.pytensorf import LayerOp
+
+
+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.TensorVariable: ...
+
+
+class LinearLayer(LayerOp):
+    __props__ = ("n_in", "n_out", "bias")
+
+
+class Linear(Layer):
+    def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
+        self.name = name if name else "Linear"
+        self.n_in = n_in
+        self.n_out = n_out
+        self.bias = bias
+
+        self.W = pt.tensor(f"{self.name}_W", shape=(n_in, self.n_out))
+
+        if self.bias:
+            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
+        inputs = [X, self.W]
+        if self.bias:
+            res += self.b
+            inputs.append(self.b)
+
+        final_st_shape = shape_to_str(res.type.shape)
+
+        ofg = LinearLayer(
+            inputs=inputs,
+            outputs=[res],
+            name=f"{self.name}[{init_st_shape} -> {final_st_shape}]",
+            n_in=self.n_in,
+            n_out=self.n_out,
+            bias=self.bias,
+        )
+        out = ofg(*inputs)
+        out.name = f"{self.name}_output"
+
+        return out
+
+
+class DropoutLayer(LayerOp):
+    __props__ = ("p",)
+
+
+class Dropout(Layer):
+    def __init__(self, name: str | None = None, p: float = 0.5, random_state: Any | None = None):
+        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(random_state))
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorLike:
+        X = pt.as_tensor(X)
+        p = pt.as_tensor(self.p, dtype=config.floatX)
+        new_rng, mask = ptr.bernoulli(p=1 - 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 - p), iff=0)],
+            name=f"{self.name}[p = {self.p}]",
+            p=self.p,
+        )(X, mask)
+        X_masked.name = f"{self.name}_output"
+
+        return X_masked
+
+
+class BatchNormLayer(LayerOp):
+    __props__ = ("n_in", "epsilon", "momentum", "affine")
+
+    def update_map(self):
+        return {1: 3, 2: 4}
+
+
+class NoRunningStatsBatchNormLayer(LayerOp):
+    __props__ = ("n_in", "epsilon", "momentum", "affine")
+
+
+class PredictionBatchNormLayer(LayerOp):
+    __props__ = ("n_in", "epsilon", "momentum", "affine")
+
+
+class BatchNorm2D(Layer):
+    def __init__(
+        self,
+        name: str | None = None,
+        n_in: int | None = None,
+        epsilon: float = 1e-5,
+        momentum: float = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+    ):
+        self.name = name if name else "BatchNorm"
+        self.n_in = n_in
+        self.epsilon = epsilon
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+
+        self.loc = None
+        self.scale = None
+
+        self.running_mean = None
+        self.running_var = None
+
+        self.initialized = False
+        self._initialize_params(None)
+
+    def _initialize_params(self, X: pt.TensorLike | None):
+        if self.initialized:
+            return
+
+        if self.n_in is None and X is None:
+            self.scale = None
+            self.loc = None
+            return
+
+        if X is not None:
+            n_in = X.type.shape[-1]
+        else:
+            n_in = self.n_in
+
+        if self.affine:
+            self.scale = pt.tensor(f"{self.name}_scale", shape=(n_in,))
+            self.loc = pt.tensor(f"{self.name}_loc", shape=(n_in,))
+            self.initialized = True
+
+        else:
+            self.initialized = True
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorLike:
+        X = pt.as_tensor(X)
+        inputs = [X]
+
+        self._initialize_params(X)
+
+        mu = X.mean(axis=0)
+        sigma_sq = X.var(axis=0)
+
+        X_normalized = (X - mu) / pt.sqrt(sigma_sq + self.epsilon)
+
+        if self.affine:
+            X_rescaled = X_normalized * self.scale + self.loc
+            inputs.extend([self.loc, self.scale])
+
+        else:
+            X_rescaled = X_normalized
+
+        if self.track_running_stats:
+            self.running_mean = mu.type(name=f"{self.name}_running_mean")
+            self.running_var = sigma_sq.type(name=f"{self.name}_running_var")
+
+            new_running_mean = self.momentum * mu + (1 - self.momentum) * self.running_mean
+            new_running_var = self.momentum * sigma_sq + (1 - self.momentum) * self.running_var
+
+            batch_norm_op = BatchNormLayer(
+                inputs=[*inputs, self.running_mean, self.running_var],
+                outputs=[X_rescaled, new_running_mean, new_running_var],
+                name=f"{self.name}",
+                n_in=self.n_in,
+                epsilon=self.epsilon,
+                momentum=self.momentum,
+                affine=self.affine,
+            )
+
+            X_transformed, self.new_running_mean, self.new_running_var = batch_norm_op(
+                *inputs, self.running_mean, self.running_var
+            )
+
+        else:
+            batch_norm_op = NoRunningStatsBatchNormLayer(
+                inputs=inputs,
+                outputs=[X_rescaled],
+                name=f"{self.name}",
+                n_in=self.n_in,
+                epsilon=self.epsilon,
+                momentum=self.momentum,
+                affine=self.affine,
+            )
+
+            X_transformed = batch_norm_op(*inputs)
+
+        X_transformed.name = f"{self.name}_output"
+
+        return X_transformed
+
+
+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__ = ["BatchNorm2D", "Concatenate", "Dropout", "Input", "Linear", "Sequential", "Squeeze"]

pytensor_ml-0.0.3/pytensor_ml/model.py

@@ -0,0 +1,161 @@
+from collections.abc import Generator
+from functools import partial
+from typing import Any, Literal, cast
+
+import numpy as np
+
+from pytensor import config
+from pytensor.compile.sharedvalue import SharedVariable
+from pytensor.graph import graph_inputs
+from pytensor.graph.basic import Constant, ancestors
+from pytensor.printing import debugprint
+from pytensor.tensor import TensorVariable
+
+from pytensor_ml.pytensorf import LayerOp, function, rewrite_for_prediction
+
+InitializationSchemes = Literal["zeros", "xavier_uniform", "xavier_normal"]
+
+
+def required_graph_inputs(tensor: TensorVariable) -> Generator[TensorVariable, None, None]:
+    return (
+        cast(TensorVariable, var)
+        for var in graph_inputs([tensor])
+        if not isinstance(var, Constant | SharedVariable)
+    )
+
+
+class Model:
+    def __init__(
+        self,
+        X: TensorVariable,
+        y: TensorVariable,
+        random_state: Any | None = None,
+        compile_kwargs: dict | None = None,
+    ):
+        self.X = X
+        self.y = y
+
+        self.rng = np.random.default_rng(random_state)
+
+        self._compile_kwargs = compile_kwargs if compile_kwargs else {}
+        self._weight_values: list[np.ndarray[float]] | None = None
+        self._non_trainable_values: list[np.ndarray[float]] | None = None
+        self._predict_fn = None
+
+    def initialize_weights(
+        self,
+        scheme: InitializationSchemes = "xavier_normal",
+        random_seed: int | str | np.random.Generator | None = None,
+    ):
+        self._weight_values = initialize_weights(self, scheme, random_seed)
+
+    def initialize_non_trainable_values(self) -> list[np.ndarray]:
+        if self.updates[0]:
+            values = [
+                self.rng.uniform(0, 1, param.type.shape).astype(param.type.dtype)
+                for param in self.updates[0]
+            ]
+            self._non_trainable_values = values
+        else:
+            self._non_trainable_values = []
+
+    @property
+    def weights(self) -> list[TensorVariable]:
+        not_trainable = {self.X}
+        for ancestor in ancestors([self.y]):
+            node = ancestor.owner
+            if node is not None and isinstance(node.op, LayerOp):
+                for input_idx in node.op.update_map().values():
+                    not_trainable.add(node.inputs[input_idx])
+
+        return [var for var in required_graph_inputs(self.y) if var not in not_trainable]
+
+    @property
+    def updates(self) -> tuple[list[TensorVariable], list[TensorVariable]]:
+        updates = {}
+        for ancestor in ancestors([self.y]):
+            node = ancestor.owner
+            if node is not None and isinstance(node.op, LayerOp):
+                for output_idx, input_idx in node.op.update_map().items():
+                    new_value = node.outputs[output_idx]
+                    old_value = node.inputs[input_idx]
+                    assert old_value.type.dtype == new_value.type.dtype
+                    updates[old_value] = new_value
+
+        if updates:
+            return list(updates.keys()), list(updates.values())
+
+        return [], []
+
+    @property
+    def weight_values(self) -> list[np.ndarray[float]]:
+        if self._weight_values is None:
+            self.initialize_weights()
+        return self._weight_values
+
+    @weight_values.setter
+    def weight_values(self, values: list[np.ndarray[float]]):
+        for i, new_value in enumerate(values):
+            self._weight_values[i][:] = new_value
+
+    @property
+    def non_trainable_values(self) -> list[np.ndarray[float]]:
+        if self._non_trainable_values is None:
+            self.initialize_non_trainable_values()
+        return self._non_trainable_values
+
+    @non_trainable_values.setter
+    def non_trainable_values(self, values: list[np.ndarray[float]]):
+        for i, new_value in enumerate(values):
+            self._non_trainable_values[i][:] = new_value
+
+    def predict(self, X_values: np.ndarray) -> np.ndarray:
+        y_pred = rewrite_for_prediction(self.y)
+        if self._predict_fn is None:
+            f = function(
+                [*self.weights, *self.updates[0], self.X],
+                y_pred,
+                **self._compile_kwargs,
+            )
+            self._predict_fn = partial(f, *self.weight_values, *self.non_trainable_values)
+
+        return cast(np.ndarray, self._predict_fn(X_values))
+
+    def __str__(self):
+        return debugprint(self.y, file="str")
+
+
+def _zero_init(shape: tuple[int], *args) -> np.ndarray:
+    return np.zeros(shape, dtype=config.floatX)
+
+
+def _xavier_uniform_init(shape: tuple[int], dtype, rng: np.random.Generator) -> np.ndarray:
+    scale = np.sqrt(6.0 / np.sum([x for x in shape if x is not None]))
+    return rng.uniform(-scale, scale, size=shape).astype(dtype)
+
+
+def _xavier_normal_init(shape: tuple[int], dtype, rng: np.random.Generator) -> np.ndarray:
+    scale = np.sqrt(2.0 / np.sum([x for x in shape if x is not None]))
+    return rng.normal(0, scale, size=shape).astype(dtype)
+
+
+initialization_factory = {
+    "zeros": _zero_init,
+    "xavier_uniform": _xavier_uniform_init,
+    "xavier_normal": _xavier_normal_init,
+}
+
+
+def initialize_weights(
+    model, scheme: InitializationSchemes, random_seed: int | str | np.random.Generator | None
+):
+    rng = np.random.default_rng(random_seed)
+
+    initial_values = []
+    for var in model.weights:
+        shape = var.type.shape
+        dtype = var.type.dtype
+        f_initialize = initialization_factory[scheme]
+        initial_values.append(f_initialize(shape, dtype, rng))
+
+    return initial_values

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

@@ -1,7 +1,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Sequence
+from functools import wraps
+from typing import Any
 
 import numpy as np
+import pytensor
+import pytensor.tensor as pt
 
 from pytensor import config
 from pytensor.compile.function.types import Function
@@ -12,6 +16,37 @@ from pytensor_ml.model import Model
 from pytensor_ml.pytensorf import function, rewrite_pregrad
 
 
+def split_output_wrapper(fn: Callable, *splits) -> Callable:
+    """
+    Wraps a function to return the weights and optimizer weights separately.
+
+    Parameters
+    ----------
+    fn: Function
+        The function to wrap.
+    splits: int
+
+    Returns
+    -------
+    Function
+        A wrapped function that returns the model weights and optimizer weights separately.
+    """
+    splits = np.cumsum(splits)
+
+    @wraps(fn)
+    def wrapped_fn(*args, **kwargs):
+        result = fn(*args, **kwargs)
+        output = []
+        cursor = 0
+        for split in splits:
+            output.append(result[cursor:split])
+            cursor = split
+        output.append(result[cursor:])
+        return output
+
+    return wrapped_fn
+
+
 class Optimizer(ABC):
     def __init__(
         self,
@@ -20,7 +55,9 @@ class Optimizer(ABC):
         ndim_out: int = 1,
         optimizer_weights: list[TensorVariable] | None = None,
         compile_kwargs: dict | None = None,
+        random_state: Any | None = None,
     ):
+        self.rng = np.random.default_rng(random_state)
         self.model = model
         self.loss_fn = loss_fn
         self.ndim_out = ndim_out
@@ -43,7 +80,10 @@ class Optimizer(ABC):
 
     def _initialize_weights(self) -> list[np.ndarray]:
         if self.optimizer_weights:
-            return [np.zeros(param.type.shape) for param in self.optimizer_weights]
+            return [
+                np.zeros(param.type.shape, dtype=param.type.dtype)
+                for param in self.optimizer_weights
+            ]
         return []
 
     def _split_weights(
@@ -78,19 +118,25 @@ class Optimizer(ABC):
 
         weights = self.model.weights
         optimizer_weights = self.optimizer_weights
+        update_inputs, update_outputs = self.model.updates
 
         loss = self.loss_fn(y, y_hat)
         loss = rewrite_pregrad(loss)
 
         all_weights = self.update_parameters(weights + optimizer_weights, loss)
+
+        weights_fn_inputs = [pytensor.In(w, mutable=True) for w in self.model.weights]
+        optimizer_weights_fn_inputs = [pytensor.In(w, mutable=True) for w in self.optimizer_weights]
+        non_trainable_fn_inputs = [pytensor.In(w, mutable=True) for w in update_inputs]
+
         fn = function(
-            [x, y, *weights, *optimizer_weights],
-            [*all_weights, loss],
+            [x, y, *weights_fn_inputs, *optimizer_weights_fn_inputs, *non_trainable_fn_inputs],
+            [*all_weights, *update_outputs, loss],
+            trust_input=True,
             **compile_kwargs,
         )
-        fn.trust_input = True
 
-        return fn
+        return split_output_wrapper(fn, len(weights), len(optimizer_weights), len(update_outputs))
 
     def step(self, x_values, y_values) -> np.ndarray:
         """
@@ -100,11 +146,19 @@ class Optimizer(ABC):
         -------
         loss
         """
-        *new_weights, loss_values = self.update_fn(
-            x_values, y_values, *self.model.weight_values, *self.optimizer_weights_values
+
+        new_weights, new_optimizer_weights, new_non_trainable_values, loss_values = self.update_fn(
+            x_values,
+            y_values,
+            *self.model.weight_values,
+            *self.optimizer_weights_values,
+            *self.model.non_trainable_values,
         )
 
-        self.model.weight_values, self.optimizer_weights_values = self._split_weights(new_weights)
+        self.model.weight_values = new_weights
+        self.model.non_trainable_values = new_non_trainable_values
+
+        self.optimizer_weights_values = new_optimizer_weights
 
         return loss_values
 
@@ -168,13 +222,57 @@ class ADAGrad(Optimizer):
 
         for param, d_loss_d_param, g2 in zip(weights, grads, optimizer_weights):
             new_g2 = g2 + d_loss_d_param**2
-            weight_update = d_loss_d_param / np.sqrt(new_g2 + self.epsilon)
+            weight_update = d_loss_d_param / pt.sqrt(new_g2 + self.epsilon)
             new_weights.append(param - self.learning_rate * weight_update)
             new_optimizer_weights.append(new_g2)
 
         return new_weights + new_optimizer_weights
 
 
+class Adadelta(Optimizer):
+    def __init__(
+        self,
+        model,
+        loss_fn,
+        *,
+        ndim_out: int = 1,
+        learning_rate: TensorLike = 1.0,
+        rho: TensorLike = 0.9,
+        epsilon: TensorLike = 1e-8,
+    ):
+        self.learning_rate = learning_rate
+        self.rho = rho
+        self.epsilon = epsilon
+
+        u_weights = [param.type() for param in model.weights]
+        v_weights = [param.type() for param in model.weights]
+        optimizer_weights = u_weights + v_weights
+        super().__init__(model, loss_fn, ndim_out=ndim_out, optimizer_weights=optimizer_weights)
+
+    def update_parameters(
+        self, weights: list[TensorVariable], loss: TensorVariable
+    ) -> list[TensorVariable]:
+        weights, optimizer_weights = self._split_weights(weights)
+        u_weights, v_weights = self._split_weights(optimizer_weights)
+
+        grads = grad(loss, weights)
+
+        new_weights = []
+        new_u_weights = []
+        new_v_weights = []
+
+        for param, d_loss_d_param, u, v in zip(weights, grads, u_weights, v_weights):
+            new_v = v * self.rho + d_loss_d_param**2 * (1 - self.rho)
+            weight_update = (sqrt(u + self.epsilon) / sqrt(new_v + self.epsilon)) * d_loss_d_param
+            new_u = u * self.rho + weight_update**2 * (1 - self.rho)
+
+            new_weights.append(param - self.learning_rate * weight_update)
+            new_u_weights.append(new_u)
+            new_v_weights.append(new_v)
+
+        return new_weights + new_u_weights + new_v_weights
+
+
 class Adam(Optimizer):
     def __init__(
         self,
@@ -232,3 +330,58 @@ class Adam(Optimizer):
             new_v_weights.append(new_v)
 
         return new_weights + new_m_weights + new_v_weights + [new_t]
+
+
+class AdamW(Optimizer):
+    def __init__(
+        self,
+        model,
+        loss_fn,
+        *,
+        ndim_out: int = 1,
+        learning_rate: TensorLike = 0.01,
+        beta1: TensorLike = 0.9,
+        beta2: TensorLike = 0.999,
+        epsilon: TensorLike = 1e-8,
+        weight_decay: TensorLike = 0.01,
+    ):
+        self.learning_rate = learning_rate
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.epsilon = epsilon
+        self.weight_decay = weight_decay
+
+        m_weights = [param.type() for param in model.weights]
+        v_weights = [param.type() for param in model.weights]
+        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)
+
+    def update_parameters(
+        self, weights: list[TensorVariable], loss: TensorVariable
+    ) -> list[TensorVariable]:
+        weights, optimizer_weights = self._split_weights(weights)
+        t = optimizer_weights.pop(-1)
+        m_weights, v_weights = self._split_weights(optimizer_weights)
+
+        grads = grad(loss, weights)
+
+        new_weights = []
+        new_m_weights = []
+        new_v_weights = []
+
+        new_t = t + 1
+        a_t = sqrt(1 - self.beta2**new_t) / (1 - self.beta1**new_t)
+
+        for param, d_loss_d_param, m, v in zip(weights, grads, m_weights, v_weights):
+            decayed_param = (1 - self.weight_decay * self.learning_rate) * param
+            weight_update = a_t * m / (sqrt(v) + self.epsilon)
+            new_weights.append(decayed_param - self.learning_rate * weight_update)
+
+            new_m = self.beta1 * m + (1 - self.beta1) * d_loss_d_param
+            new_v = self.beta2 * v + (1 - self.beta2) * d_loss_d_param**2
+            new_m_weights.append(new_m)
+            new_v_weights.append(new_v)
+
+        return new_weights + new_m_weights + new_v_weights + [new_t]

{pytensor_ml-0.0.2 → pytensor_ml-0.0.3}/pytensor_ml/pytensorf.py

@@ -5,10 +5,28 @@ 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.compile.builders import OpFromGraph
 from pytensor.graph import FunctionGraph, RewriteDatabaseQuery, rewrite_graph
 from pytensor.tensor.variable import Variable
 
-from pytensor_ml.rewriting.layers import predict_db
+
+class LayerOp(OpFromGraph):
+    # This can be removed once https://github.com/pymc-devs/pytensor/issues/1114 is fixed
+
+    __props__ = ()
+
+    def __init__(self, *args, **kwargs):
+        prop_kwargs = {key: value for key, value in kwargs.items() if key in self.__props__}
+        kwargs = {key: value for key, value in kwargs.items() if key not in self.__props__}
+
+        for key, value in prop_kwargs.items():
+            setattr(self, key, value)
+
+        super().__init__(*args, **kwargs)
+
+    def update_map(self) -> dict[int, int]:
+        """Return a mapping of output indexes to input indexes"""
+        return {}
 
 
 def atleast_list(x):
@@ -75,11 +93,13 @@ def rewrite_pregrad(graph):
 
 def rewrite_for_prediction(graph):
     """Apply rewrites to specialize a graph for forward passes (e.g. removing Dropout layers)"""
+    from pytensor_ml.rewriting.layers import predict_db
+
     if isinstance(graph, FunctionGraph):
         fgraph = graph
     else:
         outputs = [graph] if isinstance(graph, Variable) else graph
-        fgraph = FunctionGraph(outputs=outputs, clone=False)
+        fgraph = FunctionGraph(outputs=outputs, clone=True, copy_inputs=False)
 
     rewriter = predict_db.query(RewriteDatabaseQuery(include=["basic"]))
     rewriter.rewrite(fgraph)

pytensor_ml-0.0.3/pytensor_ml/rewriting/layers.py

@@ -0,0 +1,85 @@
+import pytensor.tensor as pt
+
+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 BatchNormLayer, DropoutLayer, PredictionBatchNormLayer
+
+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",
+)
+
+
+@node_rewriter([BatchNormLayer])
+def rewrite_batch_stats_to_running_average_stats(
+    fgraph: FunctionGraph, node: Apply
+) -> list[Variable] | None:
+    """
+    Replace usage of batch mean and variance with running mean and variance.
+
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        Graph being rewritten
+    node: Node
+        Node being rewritten
+
+    Returns
+    -------
+    X_normalized: Variable
+
+    """
+    X, loc, scale, running_mean, running_var = node.inputs
+
+    res = (X - running_mean) / pt.sqrt(running_var + node.op.epsilon)
+    res = loc + res * scale
+
+    batch_norm_op = PredictionBatchNormLayer(
+        inputs=[X, loc, scale, running_mean, running_var],
+        outputs=[res],
+        name=f"{node.op.name}",
+        n_in=node.op.n_in,
+        momentum=node.op.momentum,
+        epsilon=node.op.epsilon,
+        affine=node.op.affine,
+    )
+
+    X_normalized = batch_norm_op(X, loc, scale, running_mean, running_var)
+
+    return [X_normalized, None, None]
+
+
+predict_db.register(
+    "rewrite_batch_stats_to_running_average_stats",
+    rewrite_batch_stats_to_running_average_stats,
+    "basic",
+)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor-ml
-Version: 0.0.2
+Version: 0.0.3
 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.2/pytensor_ml/layers.py

@@ -1,103 +0,0 @@
-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.2/pytensor_ml/model.py

@@ -1,108 +0,0 @@
-from collections.abc import Generator
-from functools import partial
-from typing import Literal, cast
-
-import numpy as np
-
-from pytensor import config
-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"]
-
-
-def required_graph_inputs(tensor: TensorVariable) -> Generator[TensorVariable, None, None]:
-    return (
-        cast(TensorVariable, var)
-        for var in graph_inputs([tensor])
-        if not isinstance(var, Constant | SharedVariable)
-    )
-
-
-class Model:
-    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 initialize_weights(
-        self,
-        scheme: InitializationSchemes = "xavier_normal",
-        random_seed: int | str | np.random.Generator | None = None,
-    ):
-        self._weight_values = initialize_weights(self, scheme, random_seed)
-
-    @property
-    def weights(self) -> list[TensorVariable]:
-        return [var for var in required_graph_inputs(self.y) if var is not self.X]
-
-    @property
-    def weight_values(self) -> list[np.ndarray[float]]:
-        if self._weight_values is None:
-            raise ValueError("Weights have not been initialized")
-        return self._weight_values
-
-    @weight_values.setter
-    def weight_values(self, values: list[np.ndarray[float]]):
-        for i, new_value in enumerate(values):
-            self._weight_values[i][:] = new_value
-
-    def predict(self, X_values: np.ndarray) -> np.ndarray:
-        if self._predict_fn is None:
-            f = function(
-                [*self.weights, self.X],
-                rewrite_for_prediction(self.y),
-                **self._compile_kwargs,
-            )
-            self._predict_fn = partial(f, *self.weight_values)
-
-        return cast(np.ndarray, self._predict_fn(X_values))
-
-    def __str__(self):
-        return debugprint(self.y, file="str")
-
-
-def _zero_init(shape: tuple[int], *args) -> np.ndarray:
-    return np.zeros(shape, dtype=config.floatX)
-
-
-def _xavier_uniform_init(shape: tuple[int], rng: np.random.Generator) -> np.ndarray:
-    scale = np.sqrt(6.0 / np.sum([x for x in shape if x is not None]))
-    return rng.uniform(-scale, scale, size=shape).astype(config.floatX)
-
-
-def _xavier_normal_init(shape: tuple[int], rng: np.random.Generator) -> np.ndarray:
-    scale = np.sqrt(2.0 / np.sum([x for x in shape if x is not None]))
-    return rng.normal(0, scale, size=shape).astype(config.floatX)
-
-
-initialization_factory = {
-    "zeros": _zero_init,
-    "xavier_uniform": _xavier_uniform_init,
-    "xavier_normal": _xavier_normal_init,
-}
-
-
-def initialize_weights(
-    model, scheme: InitializationSchemes, random_seed: int | str | np.random.Generator | None
-):
-    if isinstance(random_seed, str):
-        random_seed = sum(map(ord, random_seed))
-    if isinstance(random_seed, int) or random_seed is None:
-        random_seed = np.random.default_rng(random_seed)
-
-    initial_values = []
-    for var in model.weights:
-        shape = var.type.shape
-        f_initialize = initialization_factory[scheme]
-        initial_values.append(f_initialize(shape, random_seed))
-
-    return initial_values

pytensor_ml-0.0.2/pytensor_ml/rewriting/layers.py

@@ -1,37 +0,0 @@
-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",
-)