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

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-30T10:40:07+0800",
+ "date": "2025-06-18T22:13:49+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "599c7aacb627660af38618e6095d8762a8d49a0a",
- "version": "0.0.1"
+ "full-revisionid": "56994236bd76d029096eb5ad6ca3767a893bb4fa",
+ "version": "0.0.3"
 }
 '''  # END VERSION_JSON
 

{pytensor_ml-0.0.1 → 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