Draw every parameter at construction, and refuse a constant scheme

pytensor_ml/layers/embedding.py

@@ -1,11 +1,8 @@
-import numpy as np
 import pytensor.tensor as pt
 
-from pytensor import config
-
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
-from pytensor_ml.state import Initializer
+from pytensor_ml.state import Initializer, XavierNormalInitializer
 
 
 class EmbeddingLayer(UnaryLayerOp):
@@ -50,8 +47,16 @@ def __init__(
         self.n_embeddings = n_embeddings
         self.n_features = n_features
 
-        W_value = np.zeros((n_embeddings, n_features), dtype=config.floatX)
-        self.W = trainable(W_value, f"{self.name}_W", initializer=weight_initializer)
+        # Drawn here, as in Linear, and for the same reason. The fallback is deliberately not declared, so
+        # a network-wide scheme still reaches the table.
+        W_initializer = (
+            XavierNormalInitializer() if weight_initializer is None else weight_initializer
+        )
+        self.W = trainable(
+            W_initializer.initial_value((n_embeddings, n_features)),
+            f"{self.name}_W",
+            initializer=weight_initializer,
+        )
 
     def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
         ids = pt.as_tensor(ids)

pytensor_ml/layers/linear.py

@@ -1,11 +1,8 @@
-import numpy as np
 import pytensor.tensor as pt
 
-from pytensor import config
-
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
-from pytensor_ml.state import Initializer, ZeroInitializer
+from pytensor_ml.state import Initializer, XavierNormalInitializer, ZeroInitializer
 
 
 def shape_to_str(shape):
@@ -47,10 +44,9 @@ class Linear(Layer):
 
     Notes
     -----
-    The weight matrix :math:`W` starts at zero, so in a stack every activation below the first layer is
-    zero and every weight matrix receives a zero gradient: an uninitialized network can fit only its
-    output bias, and predicts a constant. Call :meth:`~pytensor_ml.model.Model.initialize`, or assign a
-    value yourself, before training.
+    Both parameters are drawn when the layer is built, so a network trains without any further call.
+    :meth:`~pytensor_ml.model.Model.initialize` redraws them from a single seed, which is what makes a run
+    reproducible; :math:`W` follows its ``scheme`` there, since the draw here declares no initializer.
     """
 
     def __init__(
@@ -68,15 +64,25 @@ def __init__(
         self.n_out = n_out
         self.bias = bias
 
-        W_value = np.zeros((n_in, n_out), dtype=config.floatX)
-        self.W = trainable(W_value, f"{self.name}_W", initializer=weight_initializer)
+        # Drawn here rather than left at zero: a weight and its value are one thing, and a stack of zero
+        # weights passes no gradient below its last layer. Note what is declared and what is not --
+        # `weight_initializer` is, so it survives a network-wide scheme, while the fallback below must not
+        # be, or a scheme could never reach a weight matrix at all.
+        W_initializer = (
+            XavierNormalInitializer() if weight_initializer is None else weight_initializer
+        )
+        self.W = trainable(
+            W_initializer.initial_value((n_in, n_out)),
+            f"{self.name}_W",
+            initializer=weight_initializer,
+        )
 
         if self.bias:
-            b_value = np.zeros(n_out, dtype=config.floatX)
+            b_initializer = ZeroInitializer() if bias_initializer is None else bias_initializer
             self.b = trainable(
-                b_value,
+                b_initializer.initial_value((n_out,)),
                 f"{self.name}_b",
-                initializer=ZeroInitializer() if bias_initializer is None else bias_initializer,
+                initializer=b_initializer,
             )
 
     def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:

pytensor_ml/layers/norm.py

@@ -89,15 +89,11 @@ def _affine_parameters(
     normalizing and then rescaling by a random factor defeats the point of the layer. A caller who wants
     something else says so, and their choice becomes the declaration.
     """
-    loc = trainable(
-        np.zeros(n_in, dtype=config.floatX),
-        f"{name}_loc",
-        initializer=ZeroInitializer() if loc_initializer is None else loc_initializer,
-    )
+    resolved_loc = ZeroInitializer() if loc_initializer is None else loc_initializer
+    resolved_scale = OneInitializer() if scale_initializer is None else scale_initializer
+    loc = trainable(resolved_loc.initial_value((n_in,)), f"{name}_loc", initializer=resolved_loc)
     scale = trainable(
-        np.ones(n_in, dtype=config.floatX),
-        f"{name}_scale",
-        initializer=OneInitializer() if scale_initializer is None else scale_initializer,
+        resolved_scale.initial_value((n_in,)), f"{name}_scale", initializer=resolved_scale
     )
     return loc, scale
 

pytensor_ml/model.py

@@ -11,7 +11,11 @@
 from pytensor_ml.loss import Loss, supervised_loss
 from pytensor_ml.params import TrainableParameter
 from pytensor_ml.pytensorf import collect_trainable_params, compile_predict
-from pytensor_ml.state import InitializationSchemeLike, initialize_params
+from pytensor_ml.state import (
+    InitializationSchemeLike,
+    initialize_params,
+    require_varying_scheme,
+)
 
 
 class Model:
@@ -49,10 +53,13 @@ def initialize(
         ----------
         scheme : str or Initializer
             Initialization scheme for the weights that do not declare one: the name of a built-in
-            scheme, or an :class:`~pytensor_ml.state.Initializer` instance. Default 'xavier_normal'.
+            scheme, or an :class:`~pytensor_ml.state.Initializer` instance. A constant is refused here,
+            since identical weights leave a layer with no symmetry to break; declare one on the parameter
+            that wants it. Default 'xavier_normal'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
         """
+        require_varying_scheme(scheme)
         parameters = self.weights
         for parameter, value in zip(parameters, initialize_params(parameters, scheme, rng=seed)):
             parameter.set_value(value)

pytensor_ml/state.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 
+from pytensor import config
 from pytensor.compile.sharedvalue import SharedVariable
 
 from pytensor_ml.params import TrainableParameter
@@ -36,6 +37,20 @@ def sample(
         self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator
     ) -> np.ndarray: ...
 
+    def initial_value(self, shape: tuple[int, ...]) -> np.ndarray:
+        """
+        Draw the value a parameter of ``shape`` is born holding, at the current ``floatX``.
+
+        Uses fresh entropy: reproducibility comes from :meth:`~pytensor_ml.model.Model.initialize`, which
+        redraws every parameter from one seed and discards whatever this produced.
+
+        Parameters
+        ----------
+        shape : tuple of int
+            Shape of the parameter to draw.
+        """
+        return self.sample(shape, config.floatX, np.random.default_rng())
+
     def _sample_like(self, param: SharedVariable, rng: RandomState | None = None) -> np.ndarray:
         rng = np.random.default_rng(rng)
         value = param.get_value()
@@ -191,6 +206,34 @@ def _declared_initializer(param: SharedVariable, default: Initializer) -> Initia
     return default if declared is None else declared
 
 
+def _resolve_scheme(scheme: InitializationSchemeLike) -> Initializer:
+    """The initializer a scheme names, or the instance itself when one was passed."""
+    return scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]()
+
+
+def require_varying_scheme(scheme: InitializationSchemeLike) -> None:
+    """
+    Raise if ``scheme`` would give every weight matrix in a network the same value.
+
+    A constant is the right initializer for a bias or a norm scale and never for a whole network: identical
+    weights leave no gradient to distinguish two units in a layer, so there is no symmetry for training to
+    break. Declare it on the parameter that wants it instead.
+
+    Parameters
+    ----------
+    scheme : str or Initializer
+        The scheme about to be applied to every parameter that declares nothing.
+    """
+    if isinstance(_resolve_scheme(scheme), ZeroInitializer | OneInitializer):
+        name = scheme if isinstance(scheme, str) else type(scheme).__name__
+        raise ValueError(
+            f"{name!r} gives every weight matrix the same value, so no gradient distinguishes two units in "
+            "a layer and training cannot break the symmetry. A constant belongs on one parameter rather "
+            "than on a network: declare it where that parameter is built, as in "
+            "`Linear(..., weight_initializer=ZeroInitializer())`."
+        )
+
+
 def initialize_params(
     params: Sequence[SharedVariable],
     scheme: InitializationSchemeLike = "xavier_normal",
@@ -223,7 +266,7 @@ def initialize_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 = scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]()
+    initializer = _resolve_scheme(scheme)
     return [
         _declared_initializer(param, default=initializer)._sample_like(param, rng)
         for param in params

tests/test_layers.py

@@ -466,3 +466,69 @@ def test_a_bias_initializer_replaces_the_zero_declaration_rather_than_fighting_i
     values = dict(zip((p.name for p in params), initialize_params(params, scheme="zeros", rng=0)))
 
     assert not np.allclose(values["fc_b"], 0.0)
+
+
+@pytest.mark.parametrize("layer_name", sorted(DECLARED_BY_LAYERS), ids=sorted(DECLARED_BY_LAYERS))
+def test_a_layer_is_born_holding_a_draw_from_its_own_initializer(layer_name):
+    """Creating a parameter and giving it a value are one event, as they are in torch, flax and keras. A
+    weight left at zero passes no gradient below the last layer of a stack, so a network nobody remembered to
+    initialize trains its output bias and nothing else."""
+    build, layer_input, expected = DECLARED_BY_LAYERS[layer_name]
+    prediction = build()(layer_input)
+
+    for param in collect_trainable_params(prediction):
+        declared = expected[param.name]
+        value = param.get_value()
+        if declared is None:
+            assert len(np.unique(value)) > 1, (
+                f"{param.name} was born holding one repeated value rather than a draw"
+            )
+        else:
+            np.testing.assert_allclose(value, declared, err_msg=param.name)
+
+
+def test_a_keyword_initializer_reaches_the_value_and_not_only_the_declaration():
+    """The bias case: a keyword used to be recorded as a declaration while the value stayed zero, so it took
+    effect only if someone happened to call initialize afterwards."""
+    sentinel = 7.0
+    layer = Linear(
+        "fc",
+        n_in=4,
+        n_out=4,
+        bias_initializer=CustomInitializer(
+            lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype)
+        ),
+    )
+
+    np.testing.assert_allclose(layer.b.get_value(), sentinel)
+
+
+def test_a_non_default_scheme_still_reaches_every_weight_matrix():
+    """The trap in drawing at construction: recording the fallback as a declaration would outrank `scheme`
+    and silently disable it library-wide, while every test that initializes with the default still passed."""
+    X = pt.tensor("features", shape=(None, 4), dtype=floatX)
+    prediction = Sequential(Linear("fc1", 4, 4), Linear("fc2", 4, 4))(X)
+    params = collect_trainable_params(prediction)
+
+    values = dict(zip((p.name for p in params), initialize_params(params, scheme="ones", rng=0)))
+
+    np.testing.assert_allclose(values["fc1_W"], 1.0)
+    np.testing.assert_allclose(values["fc2_W"], 1.0)
+    np.testing.assert_allclose(values["fc1_b"], 0.0)  # its declaration still outranks the scheme
+
+
+def test_construction_draws_do_not_leak_into_a_seeded_initialize():
+    """Construction draws from fresh entropy, so two identically seeded runs must still agree: the seeded
+    initialize has to overwrite every value rather than build on what construction happened to produce."""
+
+    def seeded_values():
+        X = pt.tensor("features", shape=(None, 4), dtype=floatX)
+        prediction = Sequential(Linear("fc1", 4, 4), Linear("fc2", 4, 4))(X)
+        params = collect_trainable_params(prediction)
+        return dict(zip((p.name for p in params), initialize_params(params, rng=0)))
+
+    first, second = seeded_values(), seeded_values()
+
+    assert set(first) == set(second)
+    for name in first:
+        np.testing.assert_array_equal(first[name], second[name], err_msg=name)

tests/test_model.py

@@ -11,6 +11,7 @@
 from pytensor_ml.loss import SquaredError
 from pytensor_ml.model import Model
 from pytensor_ml.optim import sgd
+from pytensor_ml.state import OneInitializer, ZeroInitializer, initialize_params
 
 
 class TestModelPredict:
@@ -149,3 +150,48 @@ def test_compile_train_reduces_loss():
     history = [float(step(X_batch, target)) for _ in range(50)]
 
     assert history[-1] < history[0]
+
+
+@pytest.mark.parametrize(
+    "scheme",
+    ["zeros", "ones", ZeroInitializer(), OneInitializer()],
+    ids=["zeros_by_name", "ones_by_name", "zeros_instance", "ones_instance"],
+)
+def test_initialize_refuses_a_constant_scheme(scheme):
+    """A constant is the right initializer for one parameter and never for a network: identical weights leave
+    no gradient to tell two units in a layer apart. Reachable by name and by instance, so both are refused."""
+    X = pt.tensor("X", shape=(None, 8))
+    y = Sequential(Linear("fc1", 8, 4), ReLU(), Linear("fc2", 4, 2))(X)
+
+    with pytest.raises(ValueError, match="same value"):
+        Model(X, y).initialize(scheme, seed=0)
+
+
+def test_the_refusal_is_deliberately_not_symmetric_with_initialize_params():
+    """`initialize_params` takes the parameters it is handed, so "set these to zeros" is a coherent request
+    and several tests rely on a constant scheme as a distinguishable value. `Model.initialize` means the whole
+    network, where a constant is always wrong. Moving the check down to be consistent would break both this
+    contrast and about fifteen tests that use it."""
+    X = pt.tensor("X", shape=(None, 8))
+    y = Sequential(Linear("fc1", 8, 4), ReLU(), Linear("fc2", 4, 2))(X)
+    model = Model(X, y)
+
+    values = initialize_params(model.weights, scheme="zeros")
+
+    assert all(np.all(value == 0) for value in values)
+    with pytest.raises(ValueError, match="same value"):
+        model.initialize("zeros")
+
+
+def test_initialize_still_accepts_a_constant_declared_on_one_parameter():
+    """The escape hatch the refusal points at: a constant on the parameter you mean, rather than on all of
+    them. A zeroed output head is a real technique, and a declaration is how it is said."""
+    X = pt.tensor("X", shape=(None, 8))
+    first = Linear("fc1", 8, 4)
+    head = Linear("head", 4, 2, weight_initializer=ZeroInitializer())
+    y = Sequential(first, ReLU(), head)(X)
+
+    Model(X, y).initialize("xavier_normal", seed=0)
+
+    np.testing.assert_array_equal(head.W.get_value(), 0)  # its declaration outranks the scheme
+    assert np.abs(first.W.get_value()).min() > 0  # and the scheme still reached its sibling

tests/test_state.py

@@ -217,3 +217,13 @@ def test_the_normal_scheme_is_reachable_by_name():
     [value] = initialize_params([parameter], scheme="normal", rng=0)
 
     assert value.std() == pytest.approx(0.01, rel=0.05)  # the default std
+
+
+def test_initial_value_draws_at_floatx_in_the_requested_shape():
+    """What every layer relies on when it creates a parameter: the dtype matches the graph's, so nothing has
+    to remember `config.floatX` at each of the five call sites."""
+    value = XavierNormalInitializer().initial_value((6, 4))
+
+    assert value.shape == (6, 4)
+    assert str(value.dtype) == pytensor.config.floatX
+    assert len(np.unique(value)) > 1  # drawn, not filled