Replace the network-wide init scheme with per-parameter initializers

pytensor_ml/layers/attention.py

@@ -154,10 +154,9 @@ class MultiheadAttention(Layer):
     is_causal : bool, optional
         Apply a causal mask in the attention. Default is False.
     out_proj_initializer : Initializer, optional
-        How the output projection's weight is drawn, in place of whatever scheme
-        :meth:`~pytensor_ml.model.Model.initialize` is given. The three input projections are unaffected,
-        which is what a scaling applied only to the projection writing back into a residual stream needs.
-        Left to the scheme when omitted.
+        How the output projection's weight is drawn. The three input projections are unaffected, which is
+        what a scaling applied only to the projection writing back into a residual stream needs. Xavier
+        normal when omitted, as for any other weight.
     """
 
     def __init__(

pytensor_ml/layers/embedding.py

@@ -29,10 +29,10 @@ class Embedding(Layer):
     n_features : int
         Size of each embedding row.
     weight_initializer : Initializer, optional
-        How the table is drawn, in place of whatever scheme :meth:`~pytensor_ml.model.Model.initialize` is
-        given. Left to the scheme when omitted. A fan-scaled scheme puts the vocabulary size in the
-        denominator, which is correct Xavier and much tighter than the ``NormalInitializer(0.0, 0.02)`` that
-        reference implementations of GPT-2 use, so this is the keyword to reach for when matching one.
+        How the table is drawn, at construction and on every redraw. Xavier normal when omitted, which puts
+        the vocabulary size in the denominator -- correct Xavier, and much tighter than the
+        ``NormalInitializer(0.0, 0.02)`` that reference implementations of GPT-2 use, so this is the keyword
+        to reach for when matching one.
     """
 
     def __init__(
@@ -47,15 +47,14 @@ def __init__(
         self.n_embeddings = n_embeddings
         self.n_features = n_features
 
-        # 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.
+        # Drawn here, as in Linear, and for the same reason.
         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,
+            initializer=W_initializer,
         )
 
     def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:

pytensor_ml/layers/linear.py

@@ -33,11 +33,9 @@ class Linear(Layer):
     n_out : int
         Size of the output feature axis.
     bias : bool, optional
-        Add the learned shift :math:`b`, which starts at zero and stays there under a network-wide
-        initialization scheme. Default is True.
+        Add the learned shift :math:`b`, which starts at zero and is redrawn there. Default is True.
     weight_initializer : Initializer, optional
-        How :math:`W` is drawn, in place of whatever scheme :meth:`~pytensor_ml.model.Model.initialize` is
-        given. Left to the scheme when omitted, which is what makes a network-wide choice reach the weights.
+        How :math:`W` is drawn, at construction and on every redraw. Xavier normal when omitted.
     bias_initializer : Initializer, optional
         How :math:`b` is drawn. Zeros when omitted, following Keras and flax rather than torch, which draws
         the bias from :math:`\mathcal{U}(\pm 1/\sqrt{\text{fan\_in}})`.
@@ -46,7 +44,7 @@ class Linear(Layer):
     -----
     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.
+    reproducible.
     """
 
     def __init__(
@@ -65,16 +63,15 @@ def __init__(
         self.bias = bias
 
         # 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.
+        # weights passes no gradient below its last layer. The initializer is declared as well as used, so
+        # `initialize(seed=...)` redraws from the same law.
         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,
+            initializer=W_initializer,
         )
 
         if self.bias:

pytensor_ml/layers/norm.py

@@ -85,9 +85,9 @@ def _affine_parameters(
     """Build the learned shift and scale. Returns them in the ``(loc, scale)`` order that every norm
     op unpacks its inputs in, so the two cannot drift apart.
 
-    Both declare their initializer, so a network-wide scheme leaves the identity transform in place --
-    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.
+    Both declare their initializer, so a redraw returns them to the identity transform -- 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.
     """
     resolved_loc = ZeroInitializer() if loc_initializer is None else loc_initializer
     resolved_scale = OneInitializer() if scale_initializer is None else scale_initializer
@@ -167,11 +167,10 @@ class BatchNorm2D(Layer):
         0.1.
     affine : bool, optional
         Apply the learned scale :math:`\gamma` and shift :math:`\beta`, starting from the identity
-        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a network-wide initialization scheme
-        leaves in place. Default is True.
+        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a redraw returns them to. Default is True.
     scale_initializer : Initializer, optional
-        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; a scheme that
-        rescaled a normalized activation by a random factor would defeat the layer.
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; drawing a random
+        factor to rescale a normalized activation by would defeat the layer.
     loc_initializer : Initializer, optional
         How :math:`\beta` is drawn. Zeros when omitted.
     track_running_stats : bool, optional
@@ -316,11 +315,10 @@ class LayerNorm(Layer):
         Constant :math:`\epsilon` added to the variance for numerical stability. Default is 1e-5.
     affine : bool, optional
         Apply the learned scale :math:`\gamma` and shift :math:`\beta`, starting from the identity
-        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a network-wide initialization scheme
-        leaves in place. Default is True.
+        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a redraw returns them to. Default is True.
     scale_initializer : Initializer, optional
-        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; a scheme that
-        rescaled a normalized activation by a random factor would defeat the layer.
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; drawing a random
+        factor to rescale a normalized activation by would defeat the layer.
     loc_initializer : Initializer, optional
         How :math:`\beta` is drawn. Zeros when omitted.
     """

pytensor_ml/layers/transformer.py

@@ -43,9 +43,8 @@ class FeedForward(Layer):
     bias : bool, optional
         Include bias terms in both linear layers. Default is True.
     fc_out_initializer : Initializer, optional
-        How the second layer's weight is drawn, in place of whatever scheme
-        :meth:`~pytensor_ml.model.Model.initialize` is given. The hidden layer is unaffected. Left to the
-        scheme when omitted.
+        How the second layer's weight is drawn. The hidden layer is unaffected. Xavier normal when omitted,
+        as for any other weight.
     """
 
     def __init__(
@@ -125,7 +124,7 @@ class TransformerBlock(Layer):
         Constant added to the layer-norm variance for numerical stability. Default is 1e-5.
     residual_initializer : Initializer, optional
         How the two projections that write back into the residual stream are drawn -- attention's output
-        projection and the feed-forward's second layer. Their siblings are left to the scheme, which is what
+        projection and the feed-forward's second layer. Their siblings keep the default, which is what
         GPT-style initialization asks for: a residual stream accumulates one contribution per block, so those
         projections are scaled by :math:`1/\sqrt{2 n_\text{layer}}` while the rest are not. The depth is not
         known here, so the scaling belongs in the initializer the caller passes.

pytensor_ml/model.py

@@ -1,8 +1,9 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 import numpy as np
 
 from pytensor.compile import Function
+from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph.basic import Variable
 from pytensor.printing import debugprint
 from pytensor.tensor.variable import TensorVariable
@@ -11,11 +12,7 @@
 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,
-    require_varying_scheme,
-)
+from pytensor_ml.state import Initializer, initialize_params
 
 
 class Model:
@@ -39,29 +36,29 @@ def weights(self) -> list[TrainableParameter]:
 
     def initialize(
         self,
-        scheme: InitializationSchemeLike = "xavier_normal",
         seed: int | np.random.Generator | None = None,
+        initializers: Mapping[SharedVariable, Initializer] | None = None,
     ) -> "Model":
         """
-        Initialize the trainable weights in place and return self.
+        Redraw every trainable weight from its own initializer, in place, and return self.
 
-        A parameter that declares its own initializer keeps it, so ``scheme`` reaches the weights whose
-        starting value is a free choice and leaves the rest alone: a batch norm layer stays at its
-        identity transform, and biases stay at zero.
+        Each parameter was already built holding a draw, so this is what makes a run reproducible rather
+        than what makes it trainable: one seed regenerates all of them. Every layer declares how each
+        parameter it builds is drawn, so a batch norm layer redraws to its identity transform and a bias
+        to zero.
 
         Parameters
         ----------
-        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. 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.
+        initializers : dict mapping parameter to Initializer, optional
+            How to draw specific parameters, in place of what they declare. Keyed by the parameter object,
+            reached through the layer that owns it -- ``{block.ff.fc_in.b: ...}`` -- so a parameter no
+            constructor keyword exposes is still addressable.
         """
-        require_varying_scheme(scheme)
         parameters = self.weights
-        for parameter, value in zip(parameters, initialize_params(parameters, scheme, rng=seed)):
+        values = initialize_params(parameters, rng=seed, initializers=initializers)
+        for parameter, value in zip(parameters, values):
             parameter.set_value(value)
         return self
 

pytensor_ml/params.py

@@ -17,9 +17,9 @@ class TrainableParameter(TensorSharedVariable):
     Attributes
     ----------
     initializer : Initializer or None
-        The parameter's own initializer, which wins over any scheme passed to
-        :func:`~pytensor_ml.state.initialize_params`. A layer declares one when the starting value is
-        part of its definition, such as batch norm's unit scale. None defers to the caller's scheme.
+        The law this parameter's value is drawn from, which :func:`~pytensor_ml.state.initialize_params`
+        redraws it from. Every layer declares one for each parameter it builds. None means a redraw has
+        nothing to go on and raises.
     """
 
     initializer: "Initializer | None" = None
@@ -59,8 +59,8 @@ def trainable(
     Create a shared variable marked as a trainable parameter.
 
     The marker class lets graph traversal tell parameters apart from other shared state, so that an
-    optimizer updates exactly these. A parameter may also declare its own initializer, which protects a
-    meaningful starting value from being overwritten by a network-wide initialization scheme.
+    optimizer updates exactly these. A parameter also declares the initializer it is drawn from, which is
+    what lets one seed reproduce every value in a network.
 
     Parameters
     ----------
@@ -74,9 +74,8 @@ def trainable(
     strict : bool, optional
         If True, the value must exactly match the dtype.
     initializer : Initializer, optional
-        Initializer that reinitializing this parameter must use, whatever scheme the caller asks for.
-        Declare one when the starting value belongs to the layer's definition, such as a unit scale or a
-        zero bias. Default None, which defers to the caller's scheme.
+        The law to redraw this parameter from, whether that is a unit scale, a zero bias, or a fan-scaled
+        draw. Default None, which leaves the parameter with no law and raises on a redraw.
     **kwargs
         Additional arguments passed to the SharedVariable constructor.
     """