pytensor-ml 0.0.7__tar.gz → 0.1.0__tar.gz

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/.github/workflows/run_tests.yml

@@ -14,19 +14,57 @@ concurrency:
 
 jobs:
   unittest:
+    name: unittest (${{ matrix.os }}, ${{ matrix.python-version }}, ${{ matrix.test-subset.name }})
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, windows-latest]
+        os: [ubuntu-latest]
         python-version: ["3.12"]
+        # Groups are balanced on measured cold-compile-cache time, which is what a fresh runner pays.
+        # tests/test_workflow_groups.py fails if a test file is missing from every group or listed twice.
         test-subset:
-          - |
-            tests/
+          - name: attention
+            paths: |
+              tests/test_attention.py
+              tests/dispatch/
+          - name: layers
+            paths: |
+              tests/test_transformer.py
+              tests/test_layers.py
+              tests/rewriting/
+          - name: optim
+            paths: |
+              tests/optim/
+          - name: rest
+            paths: |
+              tests/test_serialize.py
+              tests/test_activations.py
+              tests/test_loss.py
+              tests/test_checkpoint.py
+              tests/test_model.py
+              tests/test_pretrained.py
+              tests/test_pytensorf.py
+              tests/test_state.py
+              tests/test_base.py
+              tests/test_collect.py
+              tests/test_util.py
+              tests/test_workflow_groups.py
+        # Windows runners are roughly twice as slow, and almost everything here is platform independent.
+        # It runs one job over the parts that are not: file IO, paths, and compiling a graph to train.
+        include:
+          - os: windows-latest
+            python-version: "3.12"
+            test-subset:
+              name: smoke
+              paths: |
+                tests/test_checkpoint.py
+                tests/test_serialize.py
+                tests/optim/test_training.py
 
     runs-on: ${{ matrix.os }}
 
     env:
-      TEST_SUBSET: ${{ matrix.test-subset }}
+      TEST_SUBSET: ${{ matrix.test-subset.paths }}
 
     defaults:
       run:
@@ -53,6 +91,16 @@ jobs:
           ~\AppData\Local\pip\Cache
         key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{ env.CACHE_NUMBER }}-${{
           hashFiles('pyproject.toml') }}
+    - name: Cache the pytensor compile directory
+      uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
+      env:
+        CACHE_NUMBER: 0
+      with:
+        path: ~/.pytensor
+        key: ${{ runner.os }}-py${{ matrix.python-version }}-pytensor-${{ env.CACHE_NUMBER }}-${{
+          matrix.test-subset.name }}-${{ hashFiles('conda_envs/pytensor_ml.yml') }}
+        restore-keys: |
+          ${{ runner.os }}-py${{ matrix.python-version }}-pytensor-${{ env.CACHE_NUMBER }}-${{ matrix.test-subset.name }}-
     - uses: conda-incubator/setup-miniconda@8ee1f361103df19b6f8c8655fd3967a8ecb162d5 # v4.0.1
       with:
         miniforge-variant: Miniforge3
@@ -75,7 +123,7 @@ jobs:
 
     - name: Run tests
       run: |
-        python -m pytest -vv --cov=pytensor_ml --cov-report=xml --no-cov-on-fail --cov-report term $TEST_SUBSET
+        python -m pytest -vv --durations=25 --cov=pytensor_ml --cov-report=xml --no-cov-on-fail --cov-report term $TEST_SUBSET
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1
       with:

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor_ml
-Version: 0.0.7
+Version: 0.1.0
 Summary: A(nother) deep learning library, built on top of Pytensor
 Project-URL: Homepage, https://github.com/pymc-devs/pytensor_ml
 Project-URL: Source, https://github.com/pymc-devs/pytensor_ml
@@ -232,7 +232,11 @@ Requires-Dist: mypy; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pytest; extra == 'dev'
 Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pyyaml; extra == 'dev'
 Requires-Dist: scikit-learn; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: matplotlib; extra == 'examples'
+Requires-Dist: tqdm; extra == 'examples'
 Description-Content-Type: text/markdown
 
 # pytensor_ml

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/conda_envs/pytensor_ml.yml

@@ -18,6 +18,7 @@ dependencies:
   - pytest-xdist
   - pytest-benchmark
   - pytest-mock
+  - pyyaml  # tests/test_workflow_groups.py reads the CI matrix
 
   # For building docs
   - sphinx

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/examples/mnist_feed_forward.ipynb

@@ -37,7 +37,7 @@
     "from sklearn.preprocessing import MinMaxScaler, OneHotEncoder\n",
     "\n",
     "from pytensor_ml.activations import LeakyReLU\n",
-    "from pytensor_ml.layers import BatchNorm2D, Dropout, Linear, Sequential\n",
+    "from pytensor_ml.layers import BatchNorm2D, Dropout, Input, Linear, Sequential\n",
     "from pytensor_ml.loss import CrossEntropy\n",
     "from pytensor_ml.model import Model\n",
     "from pytensor_ml.optim import adagrad\n",
@@ -63,8 +63,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# X_in = Input('X_in', shape=(64, ))\n",
-    "X_in = pytensor.tensor.tensor(\"X_in\", shape=(None, 64))"
+    "X_in = Input(\"X_in\", shape=(None, 64))"
    ]
   },
   {

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pyproject.toml

@@ -52,6 +52,11 @@ dev = [
     "scikit-learn",
     "mypy",
     "jax",
+    "pyyaml",
+]
+examples = [
+    "matplotlib",
+    "tqdm",
 ]
 
 [tool.hatch.version]

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pytensor_ml/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.0.7'
-__version_tuple__ = version_tuple = (0, 0, 7)
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
 
 __commit_id__ = commit_id = None

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pytensor_ml/checkpoint.py

@@ -42,19 +42,24 @@ def save_state(shared_variables: Sequence[SharedVariable], path: str | Path) ->
         Destination archive, written verbatim.
     """
     indexed = _index_by_name(shared_variables)
-    tensors = {name: _as_saveable_array(variable.get_value()) for name, variable in indexed.items()}
+    tensors = {
+        name: _as_saveable_array(variable.get_value(), variable.type.dtype)
+        for name, variable in indexed.items()
+    }
     save_file(tensors, fspath(path))
 
 
-def _as_saveable_array(value: Any) -> np.ndarray:
+def _as_saveable_array(value: Any, dtype: str) -> np.ndarray:
     """
-    Return a value as a C-contiguous numpy array safetensors will accept, preserving its rank.
+    Return a value as a C-contiguous numpy array of ``dtype`` that safetensors will accept, preserving its
+    rank.
 
     A JIT backend stores its own array type in the shared variables a compiled function updates, so a
     trained model's values are ``jax.Array`` or ``mlx.core.array``; both convert through the array
-    protocol.
+    protocol. It may also narrow them -- JAX computes in 32-bit unless ``floatX`` is ``float64`` -- so the
+    archive records the declared dtype to stay portable across precisions.
     """
-    array = np.asarray(value)
+    array = np.asarray(value, dtype=dtype)
     # Not np.ascontiguousarray: it forces ndim >= 1, reshaping a rank-0 array such as the step count.
     return array if array.flags["C_CONTIGUOUS"] else np.array(array, order="C")
 
@@ -110,11 +115,14 @@ def load_state(
     mismatches = []
     for key, variable in target_by_key.items():
         value = values[key]
-        # Through numpy: a backend array's own dtype object never compares equal to a numpy dtype.
+        # The declared dtype, not the stored value's, which a backend may have narrowed (see
+        # _as_saveable_array); set_value filters to the declaration anyway. Shape has to come from the
+        # value, whose runtime dimensions the declared type may leave unknown.
         current = np.asarray(variable.get_value(borrow=True))
-        if value.shape != current.shape or value.dtype != current.dtype:
+        declared_dtype = np.dtype(variable.type.dtype)
+        if value.shape != current.shape or value.dtype != declared_dtype:
             mismatches.append(
-                f"  {variable.name!r}: target is {current.dtype} of shape {current.shape}, "
+                f"  {variable.name!r}: target is {declared_dtype} of shape {current.shape}, "
                 f"archive has {value.dtype} of shape {value.shape}"
             )
     if mismatches:

pytensor_ml-0.1.0/pytensor_ml/layers/combinators.py

@@ -0,0 +1,34 @@
+from collections.abc import Callable
+
+import pytensor.tensor as pt
+
+
+def Input(name: str, shape: tuple[int | None, ...], dtype: str | None = None) -> pt.TensorVariable:
+    """
+    Create a named symbolic input tensor.
+
+    Parameters
+    ----------
+    name : str
+        Name of the input variable.
+    shape : tuple of int or None
+        Size of each dimension. Use None wherever the size varies between calls, such as a batch axis.
+    dtype : str, optional
+        Data type of the input. Default ``floatX``.
+    """
+    return pt.tensor(name=name, shape=shape, dtype=dtype)
+
+
+def Sequential(*layers: Callable) -> Callable:
+    """Compose layers left to right into a single callable that threads its input through each in turn."""
+
+    def forward(x: pt.TensorLike) -> pt.TensorLike:
+        for layer in layers:
+            x = layer(x)
+        return x
+
+    return forward
+
+
+Squeeze = pt.squeeze
+Concatenate = pt.concatenate

pytensor_ml-0.1.0/pytensor_ml/layers/dropout.py

@@ -0,0 +1,76 @@
+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.tensor.random.variable import RandomGeneratorSharedVariable, shared_rng
+
+from pytensor_ml.base import Layer, UnaryLayerOp
+
+
+class DropoutLayer(UnaryLayerOp):
+    __props__ = ("p",)
+
+    def build_inner_graph(self, X, mask):
+        return [pt.where(mask, ift=X / (1 - self.p), iff=0)]
+
+
+class Dropout(Layer):
+    """
+    Zero a random fraction of its input, rescaling the rest so the expected sum is unchanged.
+
+    Parameters
+    ----------
+    name : str, optional
+        Name of the layer, used to name its output and its generators. Default "Dropout".
+    p : float
+        Probability of zeroing each element. Default 0.5.
+    random_state : int, Generator, or other seed, optional
+        Seed for the layer's masks. Draws are reproducible under a given seed, including when the layer is
+        applied at several points in a network. Seeded from fresh entropy when omitted.
+
+    Attributes
+    ----------
+    generators : list of RandomGeneratorSharedVariable
+        One generator per application of the layer, in the order they were applied. Set a value on these to
+        steer or restore the masks.
+    """
+
+    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.generators: list[RandomGeneratorSharedVariable] = []
+        self._generator_source = np.random.default_rng(random_state)
+
+    def _own_generator(self) -> RandomGeneratorSharedVariable:
+        """Return a generator for one application of this layer, spawned from its seed.
+
+        Each application draws its own mask, so each needs a generator of its own: a generator read by two
+        draws has no single next state, and nothing can advance it. Spawning keeps the whole layer
+        reproducible under ``random_state`` however many times it is applied.
+        """
+        generator = shared_rng(
+            self._generator_source.spawn(1)[0], name=f"{self.name}/rng_{len(self.generators)}"
+        )
+        self.generators.append(generator)
+        return generator
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
+        X = pt.as_tensor(X)
+        p = pt.as_tensor(self.p, dtype=config.floatX)
+        _, mask = ptr.bernoulli(
+            p=1 - p, size=X.shape, rng=self._own_generator(), return_next_rng=True
+        )
+        mask = mask.astype(config.floatX)
+
+        X_masked = DropoutLayer(
+            name=f"{self.name}[p = {self.p}]",
+            p=self.p,
+        )(X, mask)
+        X_masked.name = f"{self.name}_output"
+
+        return X_masked

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pytensor_ml/layers/embedding.py

@@ -5,6 +5,7 @@ from pytensor import config
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
+from pytensor_ml.state import Initializer
 
 
 class EmbeddingLayer(UnaryLayerOp):
@@ -30,15 +31,27 @@ class Embedding(Layer):
         Number of rows in the table -- the number of distinct indices it can map.
     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.
     """
 
-    def __init__(self, name: str | None, n_embeddings: int, n_features: int):
+    def __init__(
+        self,
+        name: str | None,
+        n_embeddings: int,
+        n_features: int,
+        *,
+        weight_initializer: Initializer | None = None,
+    ):
         self.name = name if name else "Embedding"
         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")
+        self.W = trainable(W_value, f"{self.name}_W", initializer=weight_initializer)
 
     def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
         ids = pt.as_tensor(ids)

pytensor_ml-0.1.0/pytensor_ml/layers/linear.py

@@ -0,0 +1,101 @@
+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
+
+
+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 LinearLayer(UnaryLayerOp):
+    __props__ = ("n_in", "n_out", "bias")
+
+    def build_inner_graph(self, X, W, b=None):
+        res = X @ W
+        if self.bias:
+            res = res + b
+        return [res]
+
+
+class Linear(Layer):
+    r"""
+    Affine map :math:`y = x W + b`.
+
+    Parameters
+    ----------
+    name : str or None
+        Name prefix for the layer's parameters. Defaults to "Linear" when None.
+    n_in : int
+        Size of the input feature axis.
+    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.
+    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.
+    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}})`.
+
+    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.
+    """
+
+    def __init__(
+        self,
+        name: str | None,
+        n_in: int,
+        n_out: int,
+        bias: bool = True,
+        *,
+        weight_initializer: Initializer | None = None,
+        bias_initializer: Initializer | None = None,
+    ):
+        self.name = name if name else "Linear"
+        self.n_in = n_in
+        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)
+
+        if self.bias:
+            b_value = np.zeros(n_out, dtype=config.floatX)
+            self.b = trainable(
+                b_value,
+                f"{self.name}_b",
+                initializer=ZeroInitializer() if bias_initializer is None else bias_initializer,
+            )
+
+    def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
+        X = pt.as_tensor(X)
+
+        inputs = [X, self.W]
+        if self.bias:
+            inputs.append(self.b)
+
+        input_shape = shape_to_str(X.type.shape)
+        output_shape = shape_to_str((X @ self.W).type.shape)
+
+        ofg = LinearLayer(
+            name=f"{self.name}[{input_shape} -> {output_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

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pytensor_ml/layers/norm.py

@@ -5,6 +5,7 @@ from pytensor import config
 
 from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp
 from pytensor_ml.params import NonTrainableParameter, TrainableParameter, non_trainable, trainable
+from pytensor_ml.state import Initializer, OneInitializer, ZeroInitializer
 
 
 def _standardize(X, epsilon, axis, keepdims=False):
@@ -75,11 +76,29 @@ def _resolve_n_in(name: str, n_in: int | None, X: pt.TensorVariable | None) -> i
     return inferred
 
 
-def _affine_parameters(name: str, n_in: int) -> tuple[TrainableParameter, TrainableParameter]:
+def _affine_parameters(
+    name: str,
+    n_in: int,
+    loc_initializer: Initializer | None = None,
+    scale_initializer: Initializer | None = None,
+) -> tuple[TrainableParameter, TrainableParameter]:
     """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."""
-    loc = trainable(np.zeros(n_in, dtype=config.floatX), f"{name}_loc")
-    scale = trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale")
+    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.
+    """
+    loc = trainable(
+        np.zeros(n_in, dtype=config.floatX),
+        f"{name}_loc",
+        initializer=ZeroInitializer() if loc_initializer is None else loc_initializer,
+    )
+    scale = trainable(
+        np.ones(n_in, dtype=config.floatX),
+        f"{name}_scale",
+        initializer=OneInitializer() if scale_initializer is None else scale_initializer,
+    )
     return loc, scale
 
 
@@ -151,7 +170,14 @@ class BatchNorm2D(Layer):
         Weight :math:`m` of the current batch statistic in the running-average update. Default is
         0.1.
     affine : bool, optional
-        Apply the learned scale :math:`\gamma` and shift :math:`\beta`. Default is True.
+        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.
+    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.
+    loc_initializer : Initializer, optional
+        How :math:`\beta` is drawn. Zeros when omitted.
     track_running_stats : bool, optional
         Maintain running mean and variance for use at prediction time. Default is True.
 
@@ -175,6 +201,9 @@ class BatchNorm2D(Layer):
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        *,
+        scale_initializer: Initializer | None = None,
+        loc_initializer: Initializer | None = None,
     ):
         self.name = name if name else "BatchNorm"
         self.n_in = n_in
@@ -182,6 +211,8 @@ class BatchNorm2D(Layer):
         self.momentum = momentum
         self.affine = affine
         self.track_running_stats = track_running_stats
+        self._scale_initializer = scale_initializer
+        self._loc_initializer = loc_initializer
 
         self.scale: TrainableParameter | None = None
         self.loc: TrainableParameter | None = None
@@ -202,7 +233,12 @@ class BatchNorm2D(Layer):
             return
 
         if self.affine:
-            self.loc, self.scale = _affine_parameters(self.name, n_in)
+            self.loc, self.scale = _affine_parameters(
+                self.name,
+                n_in,
+                loc_initializer=self._loc_initializer,
+                scale_initializer=self._scale_initializer,
+            )
 
         if self.track_running_stats:
             zeros = np.zeros(n_in, dtype=config.floatX)
@@ -283,7 +319,14 @@ class LayerNorm(Layer):
     epsilon : float, optional
         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`. Default is True.
+        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.
+    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.
+    loc_initializer : Initializer, optional
+        How :math:`\beta` is drawn. Zeros when omitted.
     """
 
     def __init__(
@@ -292,11 +335,16 @@ class LayerNorm(Layer):
         n_in: int | None = None,
         epsilon: float = 1e-5,
         affine: bool = True,
+        *,
+        scale_initializer: Initializer | None = None,
+        loc_initializer: Initializer | None = None,
     ):
         self.name = name if name else "LayerNorm"
         self.n_in = n_in
         self.epsilon = epsilon
         self.affine = affine
+        self._scale_initializer = scale_initializer
+        self._loc_initializer = loc_initializer
 
         self.scale: TrainableParameter | None = None
         self.loc: TrainableParameter | None = None
@@ -313,7 +361,12 @@ class LayerNorm(Layer):
             return
 
         if self.affine:
-            self.loc, self.scale = _affine_parameters(self.name, n_in)
+            self.loc, self.scale = _affine_parameters(
+                self.name,
+                n_in,
+                loc_initializer=self._loc_initializer,
+                scale_initializer=self._scale_initializer,
+            )
 
         self.initialized = True
 

{pytensor_ml-0.0.7 → pytensor_ml-0.1.0}/pytensor_ml/model.py

@@ -41,11 +41,15 @@ class Model:
         """
         Initialize the trainable weights 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.
+
         Parameters
         ----------
         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'.
+            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'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
         """