Bugfixes and cleanup in state.py

pytensor_ml/model.py

@@ -8,7 +8,7 @@
 from pytensor_ml.loss import Loss, supervised_loss
 from pytensor_ml.params import TrainableParameter, collect_trainable_params
 from pytensor_ml.pytensorf import compile_predict
-from pytensor_ml.state import InitializationScheme, initialize_params
+from pytensor_ml.state import InitializationSchemeLike, initialize_params
 
 
 class Model:
@@ -26,16 +26,17 @@ def weights(self) -> list[TrainableParameter]:
 
     def initialize(
         self,
-        scheme: InitializationScheme = "xavier_normal",
+        scheme: InitializationSchemeLike = "xavier_normal",
         seed: int | np.random.Generator | None = None,
     ) -> "Model":
         """
         Initialize the trainable weights in place and return self.
 
         Parameters
         ----------
-        scheme : str
-            Initialization scheme for the weights. Default 'xavier_normal'.
+        scheme : str or Initializer
+            Initialization scheme for the weights: the name of a built-in scheme, or an
+            :class:`~pytensor_ml.state.Initializer` instance. Default 'xavier_normal'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
         """
@@ -84,8 +85,7 @@ def predict(self, X_values: np.ndarray) -> np.ndarray:
                 self.y, inputs=[self.X], compile_kwargs=self._compile_kwargs
             )
 
-        result = self._predict_fn(X_values)
-        return result if isinstance(result, np.ndarray) else np.asarray(result)
+        return np.asarray(self._predict_fn(X_values))
 
     def __str__(self):
         return debugprint(self.y, file="str")

pytensor_ml/state.py

@@ -19,20 +19,11 @@ class Initializer(ABC):
     """
     Base class for parameter initializers.
 
-    Can be used in two ways:
-    - As a class: `XavierNormalInitializer(param, rng)` - directly initializes
-    - As an instance: `init = XavierNormalInitializer(); init(param, rng)`
+    Subclasses implement :meth:`sample`. Calling an instance assigns a freshly sampled value to a
+    parameter in place, while :func:`initialize_params` calls :meth:`sample` directly and leaves the
+    assignment to its caller.
     """
 
-    def __new__(cls, param: SharedVariable | None = None, rng: RandomState | None = None):
-        # If called with a param, act as a function and initialize directly
-        if param is not None:
-            instance = object.__new__(cls)
-            cls.__init__(instance)
-            return instance(param, rng)
-        # Otherwise, return an instance for later use
-        return object.__new__(cls)
-
     def __call__(self, param: SharedVariable, rng: RandomState | None = None) -> SharedVariable:
         param.set_value(self._sample_like(param, rng))
         return param
@@ -60,29 +51,25 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
 
 class XavierUniformInitializer(Initializer):
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
-        scale = np.sqrt(6.0 / np.sum([x for x in shape if x is not None]))
+        scale = np.sqrt(6.0 / np.sum(shape))
         return rng.uniform(-scale, scale, size=shape).astype(dtype)
 
 
 class XavierNormalInitializer(Initializer):
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
-        scale = np.sqrt(2.0 / np.sum([x for x in shape if x is not None]))
+        scale = np.sqrt(2.0 / np.sum(shape))
         return rng.normal(0, scale, size=shape).astype(dtype)
 
 
 class CustomInitializer(Initializer):
-    def __new__(
-        cls,
-        sample_fn: SamplingFunction | None = None,
-        param: SharedVariable | None = None,
-        rng: RandomState | None = None,
-    ):
-        instance = object.__new__(cls)
-        if sample_fn is not None:
-            instance._sample_fn = sample_fn
-        if param is not None:
-            return instance(param, rng)
-        return instance
+    """
+    Initializer built from a sampling function.
+
+    Parameters
+    ----------
+    sample_fn : callable
+        ``(shape, dtype, rng) -> ndarray``, returning the initial value for one parameter.
+    """
 
     def __init__(self, sample_fn: SamplingFunction):
         self._sample_fn = sample_fn
@@ -91,17 +78,19 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
         return self._sample_fn(shape, dtype, rng)
 
 
-_INIT_FUNCTIONS: dict[str, type[Initializer]] = {
+_INITIALIZERS: dict[str, type[Initializer]] = {
     "zeros": ZeroInitializer,
     "xavier_uniform": XavierUniformInitializer,
     "xavier_normal": XavierNormalInitializer,
     "unit_uniform": UnitUniformInitializer,
 }
 
+InitializationSchemeLike = InitializationScheme | Initializer
+
 
 def initialize_params(
     params: Sequence[SharedVariable],
-    scheme: InitializationScheme = "xavier_normal",
+    scheme: InitializationSchemeLike = "xavier_normal",
     rng: RandomState | None = None,
 ) -> list[np.ndarray]:
     """
@@ -112,7 +101,8 @@ def initialize_params(
     params
         SharedVariables to initialize values for.
     scheme
-        Initialization scheme to use.
+        Initialization scheme to use: the name of a built-in scheme, or any :class:`Initializer`
+        instance (including a :class:`CustomInitializer` wrapping your own sampling function).
     rng
         Random number generator. If None, a new one is created.
 
@@ -121,11 +111,8 @@ def initialize_params(
     list of np.ndarray
         Initialized values matching the shapes and dtypes of params.
     """
+    # Resolve once and share: a seed handed to each _sample_like call would repeat draws across parameters.
     rng = np.random.default_rng(rng)
 
-    initializer = _INIT_FUNCTIONS[scheme]()
-    results = []
-    for var in params:
-        value = var.get_value()
-        results.append(initializer.sample(value.shape, str(value.dtype), rng))
-    return results
+    initializer = scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]()
+    return [initializer._sample_like(param, rng) for param in params]

tests/rewriting/test_layers.py

@@ -44,8 +44,16 @@ def test_remove_dropout(feature_extractor_and_rng):
     assert len([node.op for node in fg.apply_nodes if isinstance(node.op, DropoutLayer)]) == 0
 
 
-def test_rewrite_batch_stats_to_running_average_stats():
-    feature_extractor = Sequential(Linear("Layer_1", n_in=6, n_out=3), BatchNorm2D())
+@pytest.mark.parametrize(
+    "consumed_downstream", [False, True], ids=["as_graph_output", "consumed_by_a_layer"]
+)
+def test_rewrite_batch_stats_to_running_average_stats(consumed_downstream):
+    layers = [Linear("Layer_1", n_in=6, n_out=3), BatchNorm2D()]
+    if consumed_downstream:
+        # A downstream SymbolicOp type-checks its inputs more strictly than a graph output does.
+        layers.append(Linear("Layer_2", n_in=3, n_out=1))
+
+    feature_extractor = Sequential(*layers)
     X = pt.tensor("X", shape=(None, 6))
     latent = feature_extractor(X)
 

tests/test_model.py

@@ -3,114 +3,85 @@
 
 from pytensor import config
 
+import pytensor_ml.model
+
 from pytensor_ml.layers import BatchNorm2D, Linear, Sequential
 from pytensor_ml.loss import SquaredError
 from pytensor_ml.model import Model
 from pytensor_ml.optim import sgd
 
 
 class TestModelPredict:
-    def test_simple_network(self):
+    def test_matches_a_hand_computed_forward_pass(self):
         X = pt.tensor("X", shape=(None, 6))
-        mlp = Sequential(Linear("fc1", n_in=6, n_out=3), Linear("fc2", n_in=3, n_out=1))
-        y = mlp(X)
+        fc1 = Linear("fc1", n_in=6, n_out=3)
+        fc2 = Linear("fc2", n_in=3, n_out=1)
+        model = Model(X, Sequential(fc1, fc2)(X)).initialize(seed=42)
 
-        model = Model(X, y)
-        model.initialize(seed=42)
+        X_test = np.random.default_rng(0).normal(size=(10, 6)).astype(config.floatX)
+        hidden = X_test @ fc1.W.get_value() + fc1.b.get_value()
+        expected = hidden @ fc2.W.get_value() + fc2.b.get_value()
 
-        X_test = np.random.randn(10, 6).astype(config.floatX)
         result = model.predict(X_test)
 
         assert result.shape == (10, 1)
         assert result.dtype == config.floatX
+        np.testing.assert_allclose(result, expected, rtol=1e-5)
 
-    def test_with_batchnorm(self):
-        X = pt.tensor("X", shape=(None, 6))
-        network = Sequential(
-            Linear("fc1", n_in=6, n_out=3),
-            BatchNorm2D("bn1", n_in=3),
-            Linear("fc2", n_in=3, n_out=1),
-        )
-        y = network(X)
-
-        model = Model(X, y)
-        model.initialize(seed=42)
-
-        X_test = np.random.randn(10, 6).astype(config.floatX)
-        result = model.predict(X_test)
-
-        assert result.shape == (10, 1)
-
-    def test_predict_uses_running_stats(self):
-        """Verify that predict uses running stats (via rewrite) not batch stats."""
+    def test_normalizes_with_running_stats_not_batch_stats(self):
         X = pt.tensor("X", shape=(None, 4))
         fc1 = Linear("fc1", n_in=4, n_out=4)
         bn = BatchNorm2D("bn1", n_in=4)
-        network = Sequential(fc1, bn)
-        y = network(X)
-
-        model = Model(X, y)
-        model.initialize(seed=42)
+        model = Model(X, Sequential(fc1, bn)(X)).initialize(seed=42)
 
-        # Set specific running stats
         bn.running_mean.set_value(np.array([1.0, 2.0, 3.0, 4.0], dtype=config.floatX))
-        bn.running_var.set_value(np.array([1.0, 1.0, 1.0, 1.0], dtype=config.floatX))
+        bn.running_var.set_value(np.ones(4, dtype=config.floatX))
 
-        # Predict with two different batches - should give same normalization
-        # if using running stats (batch stats would differ)
-        X1 = np.random.randn(5, 4).astype(config.floatX)
-        X2 = np.random.randn(20, 4).astype(config.floatX)
-
-        # Get FC output before normalization for both batches
-        fc_weight = fc1.W.get_value()
-        fc_bias = fc1.b.get_value()
-        fc_out1 = X1 @ fc_weight + fc_bias
-        fc_out2 = X2 @ fc_weight + fc_bias
-
-        # Expected: normalized using running stats
-        scale = bn.scale.get_value()
-        loc = bn.loc.get_value()
-        running_mean = bn.running_mean.get_value()
-        running_var = bn.running_var.get_value()
-
-        expected1 = (fc_out1 - running_mean) / np.sqrt(running_var + bn.epsilon) * scale + loc
-        expected2 = (fc_out2 - running_mean) / np.sqrt(running_var + bn.epsilon) * scale + loc
+        rng = np.random.default_rng(0)
+        # Two batch sizes: batch statistics would differ between them, running statistics cannot.
+        for n_rows in (5, 20):
+            X_test = rng.normal(size=(n_rows, 4)).astype(config.floatX)
+            fc_out = X_test @ fc1.W.get_value() + fc1.b.get_value()
+            standardized = (fc_out - bn.running_mean.get_value()) / np.sqrt(
+                bn.running_var.get_value() + bn.epsilon
+            )
+            expected = standardized * bn.scale.get_value() + bn.loc.get_value()
+
+            np.testing.assert_allclose(model.predict(X_test), expected, rtol=1e-5)
+
+    def test_compiles_once_and_reuses_the_function(self, monkeypatch):
+        X = pt.tensor("X", shape=(None, 4))
+        model = Model(X, Linear("fc1", n_in=4, n_out=2)(X)).initialize(seed=42)
 
-        result1 = model.predict(X1)
-        result2 = model.predict(X2)
+        uncounted_compile_predict = pytensor_ml.model.compile_predict
+        compile_count = 0
 
-        np.testing.assert_allclose(result1, expected1, rtol=1e-5)
-        np.testing.assert_allclose(result2, expected2, rtol=1e-5)
+        def counting_compile_predict(*args, **kwargs):
+            nonlocal compile_count
+            compile_count += 1
+            return uncounted_compile_predict(*args, **kwargs)
 
-    def test_compile_train_reduces_loss(self):
-        X = pt.tensor("X", shape=(None, 4))
-        y = Sequential(Linear("fc1", n_in=4, n_out=8), Linear("fc2", n_in=8, n_out=1))(X)
-        model = Model(X, y).initialize(seed=0)
+        monkeypatch.setattr(pytensor_ml.model, "compile_predict", counting_compile_predict)
 
-        step = model.compile_train(sgd(learning_rate=1e-2), SquaredError(), ndim_out=2)
+        X_test = np.random.default_rng(0).normal(size=(5, 4)).astype(config.floatX)
+        first = model.predict(X_test)
+        second = model.predict(X_test)
 
-        rng = np.random.default_rng(0)
-        X_batch = rng.normal(size=(64, 4)).astype(config.floatX)
-        target = rng.normal(size=(64, 1)).astype(config.floatX)
-        history = [float(step(X_batch, target)) for _ in range(50)]
-        assert history[-1] < history[0]
+        assert compile_count == 1
+        np.testing.assert_array_equal(first, second)
 
-    def test_predict_caches_function(self):
-        """Verify that predict function is compiled once and reused."""
-        X = pt.tensor("X", shape=(None, 4))
-        y = Linear("fc1", n_in=4, n_out=2)(X)
 
-        model = Model(X, y)
-        model.initialize(seed=42)
+def test_compile_train_reduces_loss():
+    X = pt.tensor("X", shape=(None, 4))
+    y = Sequential(Linear("fc1", n_in=4, n_out=8), Linear("fc2", n_in=8, n_out=1))(X)
+    model = Model(X, y).initialize(seed=0)
 
-        assert model._predict_fn is None
+    step = model.compile_train(sgd(learning_rate=1e-2), SquaredError(), ndim_out=2)
 
-        X_test = np.random.randn(5, 4).astype(config.floatX)
-        model.predict(X_test)
+    rng = np.random.default_rng(0)
+    X_batch = rng.normal(size=(64, 4)).astype(config.floatX)
+    target = rng.normal(size=(64, 1)).astype(config.floatX)
 
-        assert model._predict_fn is not None
-        fn_id = id(model._predict_fn)
+    history = [float(step(X_batch, target)) for _ in range(50)]
 
-        # Second call should reuse same function
-        model.predict(X_test)
-        assert id(model._predict_fn) == fn_id
+    assert history[-1] < history[0]