From 4308bf4aadb405257514315f4f085c0322f9e4f9 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 9 Aug 2026 22:50:43 -0500 Subject: [PATCH 1/3] Build Dropout generators with pytensor's shared_rng Passing a Generator to `shared` already returns a RandomGeneratorSharedVariable, so this changes which constructor is written rather than what is built. --- pytensor_ml/layers/dropout.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pytensor_ml/layers/dropout.py b/pytensor_ml/layers/dropout.py index 0b176e7..05701a2 100644 --- a/pytensor_ml/layers/dropout.py +++ b/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) From 4bb00f1d0a7182e878bca3d1b8f2e85e4711a775 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 9 Aug 2026 23:46:34 -0500 Subject: [PATCH 2/3] Fold the receptive field into the Xavier fans Summing the shape is the fan computation only for a matrix, so a 3x3 kernel came out around 2.7x too wide; the leading dimension is the fan-in here, matching this library's (n_in, n_out) layout rather than torch's transpose of it. --- pytensor_ml/state.py | 43 ++++++++++++++++++++++++++++++++-- tests/test_state.py | 55 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py index 00e4958..3b4d2b4 100644 --- a/pytensor_ml/state.py +++ b/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) diff --git a/tests/test_state.py b/tests/test_state.py index 80766f3..775ed5c 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -4,6 +4,7 @@ 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 ( @@ -11,7 +12,10 @@ 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)) From acdc8cc06ae6f716e62b62ad831dd55881a4e296 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 9 Aug 2026 23:46:49 -0500 Subject: [PATCH 3/3] Declare the examples' matplotlib and tqdm dependencies examples/mnist_feed_forward.ipynb imports both and neither was listed anywhere, so the tracked example could not be run from a declared install. --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 8920739..275bc87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,10 @@ dev = [ "jax", "pyyaml", ] +examples = [ + "matplotlib", + "tqdm", +] [tool.hatch.version] source = "vcs"