Fold the receptive field into the Xavier fans

pyproject.toml

@@ -54,6 +54,10 @@ dev = [
     "jax",
     "pyyaml",
 ]
+examples = [
+    "matplotlib",
+    "tqdm",
+]
 
 [tool.hatch.version]
 source = "vcs"

pytensor_ml/layers/dropout.py

@@ -5,8 +5,7 @@
 import pytensor.tensor.random as ptr
 
 from pytensor import config
-from pytensor.compile.sharedvalue import shared
-from pytensor.tensor.random.variable import RandomGeneratorSharedVariable
+from pytensor.tensor.random.variable import RandomGeneratorSharedVariable, shared_rng
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 
@@ -54,7 +53,7 @@ def _own_generator(self) -> RandomGeneratorSharedVariable:
         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(
+        generator = shared_rng(
             self._generator_source.spawn(1)[0], name=f"{self.name}/rng_{len(self.generators)}"
         )
         self.generators.append(generator)

pytensor_ml/state.py

@@ -55,15 +55,54 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
         return rng.uniform(0.0, 1.0, size=shape).astype(dtype)
 
 
+def fans(shape: tuple[int, ...]) -> tuple[int, int]:
+    r"""
+    Return the number of units feeding into and out of one position of a parameter of ``shape``.
+
+    Weights here are laid out input dimension first, as :class:`~pytensor_ml.layers.Linear` builds
+    ``(n_in, n_out)`` for ``X @ W`` and :class:`~pytensor_ml.layers.Embedding` builds
+    ``(vocabulary, features)``, so the leading dimension is the fan-in and the second the fan-out. This is
+    the transpose of torch's convention, where the output dimension leads. Any dimension past the second is
+    a receptive field: every input reaches an output at each of its offsets, so both fans carry a factor of
+    :math:`\prod \text{kernel}`.
+
+    Only the sum of the two matters to a Xavier draw, which is why the orientation is invisible there and
+    load-bearing for anything scaling by fan-in alone.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        Shape of the parameter, with at least two dimensions.
+
+    Returns
+    -------
+    fan_in : int
+        Units feeding one output position.
+    fan_out : int
+        Output positions one input feeds.
+    """
+    if len(shape) < 2:
+        raise ValueError(
+            f"A fan-scaled initializer needs a parameter of at least two dimensions to size its draws, but "
+            f"got shape {shape}. A bias or a norm scale has no fans; give it an initializer of its own -- "
+            "`trainable(value, name, initializer=ZeroInitializer())` -- or initialize it with the 'zeros' "
+            "or 'ones' scheme."
+        )
+    receptive_field = int(np.prod(shape[2:]))
+    return shape[0] * receptive_field, shape[1] * receptive_field
+
+
 class XavierUniformInitializer(Initializer):
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
-        scale = np.sqrt(6.0 / np.sum(shape))
+        fan_in, fan_out = fans(shape)
+        scale = np.sqrt(6.0 / (fan_in + fan_out))
         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(shape))
+        fan_in, fan_out = fans(shape)
+        scale = np.sqrt(2.0 / (fan_in + fan_out))
         return rng.normal(0, scale, size=shape).astype(dtype)
 
 

tests/test_state.py

@@ -4,14 +4,18 @@
 import pytensor
 import pytest
 
+from pytensor_ml.layers import Linear
 from pytensor_ml.params import trainable
 from pytensor_ml.pytensorf import collect_trainable_params
 from pytensor_ml.state import (
     _INITIALIZERS,
     CustomInitializer,
     InitializationScheme,
     OneInitializer,
+    XavierNormalInitializer,
+    XavierUniformInitializer,
     ZeroInitializer,
+    fans,
     initialize_params,
 )
 
@@ -128,3 +132,54 @@ def test_calling_an_initializer_assigns_the_parameter_in_place():
 
     np.testing.assert_array_equal(param.get_value(), 0)
     assert returned is param
+
+
+def test_a_convolution_kernel_folds_its_receptive_field_into_both_fans():
+    """Summing the shape is the fan computation only for a matrix. For an ``(in, out, kH, kW)`` kernel every
+    input channel reaches an output at each of the kH*kW offsets, so leaving the receptive field out of the
+    fans overstates the spread: 0.258 where the correct scale for this shape is 0.096."""
+    kernel_shape = (8, 16, 3, 3)  # asymmetric, so the orientation of the two fans is pinned as well
+    fan_in, fan_out = fans(kernel_shape)
+    assert (fan_in, fan_out) == (8 * 9, 16 * 9)
+
+    value = XavierNormalInitializer().sample(kernel_shape, "float64", np.random.default_rng(0))
+
+    assert value.std() == pytest.approx(np.sqrt(2.0 / (fan_in + fan_out)), rel=0.05)
+
+
+@pytest.mark.parametrize(
+    "shape", [(768, 768), (50257, 768), (4, 7)], ids=["square", "embedding", "small"]
+)
+def test_a_weight_matrix_draws_exactly_as_it_did_before(shape):
+    """No parameter that already existed may move: a matrix has no dimensions past the second, so the sum of
+    its fans is the sum of its shape, which is what the draw was scaled by before."""
+    fan_in, fan_out = fans(shape)
+    assert fan_in + fan_out == sum(shape)
+
+    with_fans = XavierNormalInitializer().sample(shape, "float64", np.random.default_rng(7))
+    with_shape_sum = np.random.default_rng(7).normal(0, np.sqrt(2.0 / sum(shape)), size=shape)
+
+    np.testing.assert_array_equal(with_fans, with_shape_sum)
+
+
+def test_the_fan_in_is_the_dimension_the_layers_treat_as_input():
+    """Weights here are ``(n_in, n_out)``, the transpose of torch's layout, so the leading dimension is the
+    fan-in. Xavier only reads the sum and cannot tell the difference; anything scaling by fan-in alone can."""
+    layer = Linear("fc", n_in=4, n_out=7)
+
+    fan_in, fan_out = fans(layer.W.get_value().shape)
+
+    assert (fan_in, fan_out) == (4, 7)
+
+
+@pytest.mark.parametrize(
+    "initializer",
+    [XavierNormalInitializer(), XavierUniformInitializer()],
+    ids=["normal", "uniform"],
+)
+def test_a_fan_scaled_initializer_rejects_a_parameter_with_no_fans(initializer):
+    """A bias or a norm scale has no fan-in and fan-out, and scaling by the length of a vector is a number
+    with no meaning behind it. Such parameters declare their own initializer; reaching one with a scheme
+    instead should say so rather than draw something arbitrary."""
+    with pytest.raises(ValueError, match="at least two dimensions"):
+        initializer.sample((768,), "float64", np.random.default_rng(0))