From 76c05866e2c0c6b28c861e214f5def390bba1315 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 14:48:40 -0500 Subject: [PATCH 01/10] Split the layers CI group into three --- .github/workflows/run_tests.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 76e1d24..115b9b8 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -35,13 +35,17 @@ jobs: paths: | tests/test_transformer.py tests/test_layers.py + - name: recurrent + paths: | tests/test_recurrent.py + tests/rewriting/ + - name: convolution + paths: | tests/test_conv.py tests/test_conv_transpose.py tests/test_padding.py tests/test_pool.py tests/test_upsample.py - tests/rewriting/ - name: optim paths: | tests/optim/ From bf05e83a16e0b241a498d82b0fe1b4c3a6bf6794 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 14:48:41 -0500 Subject: [PATCH 02/10] Zero masked padding before the recurrent step --- pytensor_ml/layers/recurrent.py | 5 ++++- tests/test_recurrent.py | 20 ++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/pytensor_ml/layers/recurrent.py b/pytensor_ml/layers/recurrent.py index c37ef5f..c89c5c9 100644 --- a/pytensor_ml/layers/recurrent.py +++ b/pytensor_ml/layers/recurrent.py @@ -229,7 +229,10 @@ def _masked_step( ) -> tuple[TensorVariable, ...]: """Take the step, then keep it only where ``mask_t`` says this step is real.""" # The step runs either way -- scan's carried states are fixed-shape buffers, so there is no - # skipping a subset of the batch, only discarding what it computed for them. + # skipping a subset of the batch, only discarding what it computed for them. Padding is zeroed + # first because the switch below still sends its zero cotangent through the step's derivative + # at the pad, and a non-finite pad turns that product into NaN. + x_t = pt.switch(pt.shape_padright(mask_t, x_t.ndim - mask_t.ndim), x_t, 0) stepped = self.cell.step(x_t, *state) # One trailing axis per feature axis of the state, so the mask lines its batch axes up against # the state's however many feature axes the cell chose to carry. diff --git a/tests/test_recurrent.py b/tests/test_recurrent.py index 35e25f8..f84d4bb 100644 --- a/tests/test_recurrent.py +++ b/tests/test_recurrent.py @@ -1210,6 +1210,26 @@ def test_a_mask_keeps_padding_out_of_the_gradient(rng): ) +def test_non_finite_padding_leaves_the_gradient_finite(rng): + """Padding is only a placeholder, so a batch padded with NaN has to train exactly like one padded + with zeros. The masked step's discarded branch still meets the padding on the backward pass, where + a zero cotangent times the step's derivative at NaN is NaN.""" + X = pt.tensor("X", shape=(None, None, 4)) + mask = pt.tensor("mask", shape=(None, None), dtype=bool) + layer = RNN("rnn", n_in=4, n_hidden=3, reverse=True) + draw_parameters(layer, rng) + parameters = [layer.cell.W_ih, layer.cell.W_hh, layer.cell.b] + gradients = pytensor.function([X, mask], pt.grad(layer(X, mask=mask).sum(), parameters)) + + zero_padded, mask_np = pad_to([rng.normal(size=(3, 4)).astype(floatX)], padded_length=6) + nan_padded = np.where(mask_np[..., None], zero_padded, np.nan).astype(floatX) + + for nan_gradient, zero_gradient in zip( + gradients(nan_padded, mask_np), gradients(zero_padded, mask_np) + ): + np.testing.assert_array_equal(nan_gradient, zero_gradient) + + def test_a_mask_may_skip_a_step_in_the_middle(rng): """A mask says which steps count, not how many, so it can drop one from the middle -- which a per-example length cannot express. The recurrence carries on as if that step were not there.""" From d6bacd2867684fd657a1bd5a8c017eab76a6eb57 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 14:48:41 -0500 Subject: [PATCH 03/10] Keep the input dtype through bilinear upsampling --- pytensor_ml/layers/conv.py | 22 ++++++++++++++++------ tests/test_upsample.py | 10 ++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/pytensor_ml/layers/conv.py b/pytensor_ml/layers/conv.py index a41c654..fc26ffc 100644 --- a/pytensor_ml/layers/conv.py +++ b/pytensor_ml/layers/conv.py @@ -10,6 +10,7 @@ from pytensor.gradient import disconnected_type from pytensor.graph.basic import Apply, Variable from pytensor.graph.op import Op +from pytensor.scalar import upcast from pytensor.tensor.basic import get_scalar_constant_value from pytensor.tensor.pad import PadMode from pytensor.tensor.variable import TensorVariable @@ -1305,16 +1306,19 @@ def _nearest_source_indices( def _linear_source_coordinates( - in_extent: int | TensorVariable, out_extent: int | TensorVariable, align_corners: bool + in_extent: int | TensorVariable, + out_extent: int | TensorVariable, + align_corners: bool, + dtype: str, ) -> TensorVariable: """Fractional input position each output position samples, in input coordinates. Without ``align_corners`` the outermost samples fall outside the input, which is what the clamp handles.""" # Cast the extents before dividing: they arrive as int64, and an int64 division is float64, - # which would carry the whole interpolation off floatX and back in the layer's output. - in_span = pt.cast(in_extent, config.floatX) - out_span = pt.cast(out_extent, config.floatX) + # which would carry the whole interpolation off ``dtype`` and back in the layer's output. + in_span = pt.cast(in_extent, dtype) + out_span = pt.cast(out_extent, dtype) - positions = pt.arange(out_extent, dtype=config.floatX) + positions = pt.arange(out_extent, dtype=dtype) if align_corners: # The maximum keeps a single-element output spreading over nothing rather than dividing by it. source = positions * ((in_span - 1) / pt.maximum(out_span - 1, 1.0)) @@ -1335,7 +1339,11 @@ def _resample_axis( if mode == "nearest": return pt.take(X, _nearest_source_indices(in_extent, out_extent), axis=axis) - source = _linear_source_coordinates(in_extent, out_extent, align_corners) + # Positions are computed at float32 or wider, because float16 cannot count past 2048 exactly, and + # the weights are then cast to the input's dtype so the output keeps it. + is_float = X.dtype.startswith("float") + coordinate_dtype = upcast(X.dtype, "float32") if is_float else config.floatX + source = _linear_source_coordinates(in_extent, out_extent, align_corners, coordinate_dtype) lower_position = pt.floor(source) lower_index = pt.cast(lower_position, "int64") upper_index = pt.minimum(lower_index + 1, in_extent - 1) @@ -1343,6 +1351,8 @@ def _resample_axis( # The weight varies along `axis` alone, so it broadcasts against every other axis of the gather. along_axis = (np.newaxis,) * axis + (slice(None),) + (np.newaxis,) * (X.ndim - axis - 1) weight = (source - lower_position)[along_axis] + if is_float: + weight = weight.astype(X.dtype) lower = pt.take(X, lower_index, axis=axis) upper = pt.take(X, upper_index, axis=axis) diff --git a/tests/test_upsample.py b/tests/test_upsample.py index 6aac0c6..fdcadd1 100644 --- a/tests/test_upsample.py +++ b/tests/test_upsample.py @@ -176,6 +176,16 @@ def test_resampling_keeps_the_input_dtype(dtype, mode): assert Upsample2D(scale_factor=2, mode=mode)(X).dtype == dtype +@pytest.mark.parametrize("dtype", ["float16", "float32"]) +def test_bilinear_keeps_an_input_narrower_than_floatx(dtype): + """The interpolation weights are built from integer extents, and casting them to floatX would carry + a float32 or float16 input up to float64 under the default config.""" + with pytensor.config.change_flags(floatX="float64"): + X = pt.tensor("X", shape=(None, 3, 4, 2), dtype=dtype) + + assert Upsample2D(scale_factor=2, mode="bilinear")(X).dtype == dtype + + # Every way of asking for a resampling that means nothing. `mode` is spelled for its rank, so the # name that works on one class is wrong on the other rather than a harmless alias. REJECTED_CONFIGURATIONS = [ From f4ead3a4d365c43fe5504165bc5e0457d3d67f7e Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:15:45 -0500 Subject: [PATCH 04/10] Trim the recurrent tests to one compile per checked behavior --- tests/test_recurrent.py | 818 ++++++++-------------------------------- 1 file changed, 152 insertions(+), 666 deletions(-) diff --git a/tests/test_recurrent.py b/tests/test_recurrent.py index f84d4bb..cfd3eda 100644 --- a/tests/test_recurrent.py +++ b/tests/test_recurrent.py @@ -3,25 +3,19 @@ import pytensor.tensor as pt import pytest -from pytensor_ml.activations import Activation, ReLU, Tanh +from pytensor_ml.activations import Activation, ReLU from pytensor_ml.layers import ( GRU, LSTM, RNN, Bidirectional, - ElmanCell, - GRUCell, - Input, - Linear, Recurrent, RecurrentCell, ) -from pytensor_ml.loss import SquaredError -from pytensor_ml.model import Model -from pytensor_ml.optim import adam from pytensor_ml.params import trainable from pytensor_ml.pytensorf import collect_trainable_params -from pytensor_ml.state import OneInitializer, ZeroInitializer +from pytensor_ml.state import OneInitializer +from tests.conftest import constant floatX = pytensor.config.floatX @@ -55,14 +49,9 @@ def draw_parameters(layer, rng): return W_ih, b, W_hh -@pytest.mark.parametrize( - "activation, reference", - [(Tanh(), np.tanh), (ReLU(), lambda x: np.maximum(x, 0.0))], - ids=["tanh", "relu"], -) -def test_matches_a_step_by_step_reference(activation, reference, rng): +def test_matches_a_step_by_step_reference(rng): X = pt.tensor("X", shape=(None, None, 4)) - layer = RNN("rnn", n_in=4, n_hidden=3, activation=activation) + layer = RNN("rnn", n_in=4, n_hidden=3) out = layer(X) assert out.type.shape == (None, None, 3) @@ -70,60 +59,9 @@ def test_matches_a_step_by_step_reference(activation, reference, rng): X_np = rng.normal(size=(5, 7, 4)).astype(floatX) np.testing.assert_allclose( - out.eval({X: X_np}), unrolled(X_np, W_ih, b, W_hh, reference), atol=ATOL - ) - - -def test_starts_from_a_given_state(rng): - """The state a caller hands in has to reach the first step, not just sit in the graph -- a zeros - default that quietly ignored it would agree with the reference on every other test here.""" - X = pt.tensor("X", shape=(None, None, 4)) - h0 = pt.tensor("h0", shape=(None, 3)) - layer = RNN("rnn", n_in=4, n_hidden=3) - out = layer(X, h0) - - W_ih, b, W_hh = draw_parameters(layer, rng) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - h0_np = rng.normal(size=(5, 3)).astype(floatX) - - np.testing.assert_allclose( - out.eval({X: X_np, h0: h0_np}), - unrolled(X_np, W_ih, b, W_hh, np.tanh, h0=h0_np), - atol=ATOL, - ) - - -def test_the_recurrent_weight_first_acts_on_the_second_step(rng): - """The recurrence is real and runs forward in time. Starting from a zero state, the first output does - not touch the recurrent weight at all and the second does, which is what separates a scan from an - input projection applied position by position.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = RNN("rnn", n_in=4, n_hidden=3) - out = layer(X) - draw_parameters(layer, rng) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - - first, second = ( - pytensor.grad(out[..., step, :].sum(), layer.cell.W_hh).eval({X: X_np}) for step in (0, 1) + out.eval({X: X_np}), unrolled(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL ) - np.testing.assert_allclose(first, np.zeros((3, 3)), atol=ATOL) - assert np.abs(second).max() > 1e-3 - - -def test_every_parameter_is_reachable_through_the_scan(rng): - """The recurrent weight enters the graph as a non-sequence of the scan rather than as a plain input, so - a collector that stopped at the scan node would train the projection and leave the recurrence frozen.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = RNN("rnn", n_in=4, n_hidden=3) - out = layer(X) - - assert set(collect_trainable_params(out)) == { - layer.cell.W_ih, - layer.cell.b, - layer.cell.W_hh, - } - def test_an_activation_brings_its_own_parameters_into_the_recurrence(): """The step closes over whatever the activation holds, and scan lifts it in. A strict scan would reject @@ -145,76 +83,77 @@ def __call__(self, x): assert activation.slope in collect_trainable_params(out) -def test_trains_end_to_end(rng): - """Gradients survive the round trip through the scan and the training machinery moves the parameters.""" - X = Input("X", shape=(None, 6, 4)) - y = Linear("head", n_in=5, n_out=1)(RNN("rnn", n_in=4, n_hidden=5)(X)[..., -1, :]) - model = Model(y).initialize(seed=1) - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) - - X_np = rng.normal(size=(32, 6, 4)).astype(floatX) - y_np = X_np.sum(axis=(1, 2))[:, None].astype(floatX) - - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 - - -def test_the_recurrent_weight_is_drawn_orthogonal_by_default(): +@pytest.mark.parametrize( + "layer_type, n_gates", [(RNN, 1), (GRU, 3), (LSTM, 4)], ids=["rnn", "gru", "lstm"] +) +def test_the_recurrent_weight_is_drawn_orthogonal_by_default(layer_type, n_gates): """Applied once per step, so its singular values compound: at one they leave the state alone however long the sequence, and spread around one they explode the gradient along some directions while - vanishing it along others. The input weight keeps the usual fan-scaled draw, checked structurally as - well as by spread -- on a square matrix the two draws have the same entry standard deviation, so - spread alone would not notice it picking up the recurrent default.""" - layer = RNN("rnn", n_in=16, n_hidden=64) + vanishing it along others. One draw covers every gate, as in keras, so the check is on the whole + wide matrix rather than on each gate's block. The input weight keeps the usual fan-scaled draw, + checked structurally as well as by spread -- on a square matrix the two draws have the same entry + standard deviation, so spread alone would not notice it picking up the recurrent default.""" + layer = layer_type("layer", n_in=16, n_hidden=32) W_hh = layer.cell.W_hh.get_value() - np.testing.assert_allclose(W_hh.T @ W_hh, np.eye(64), atol=ATOL) + assert W_hh.shape == (32, n_gates * 32) + np.testing.assert_allclose(W_hh @ W_hh.T, np.eye(32), atol=ATOL) W_ih = layer.cell.W_ih.get_value() assert np.abs(W_ih @ W_ih.T - np.eye(16)).max() > 0.1 - assert W_ih.std() == pytest.approx(np.sqrt(2.0 / 80), rel=0.1) + # Both fans count the stacked axis, so the spread would be wrong if the draw saw one gate's shape. + assert W_ih.std() == pytest.approx(np.sqrt(2.0 / (16 + n_gates * 32)), rel=0.1) -@pytest.mark.parametrize("bias", [True, False], ids=["bias", "no_bias"]) -def test_the_bias_is_optional(bias, rng): +def test_the_bias_is_optional(rng): """Dropping the bias has to drop the parameter as well as the term. Leaving an unused one behind would hand the optimizer moment state to carry for a weight that never moves, and nothing else here builds the layer without it.""" X = pt.tensor("X", shape=(None, None, 4)) - layer = RNN("rnn", n_in=4, n_hidden=3, bias=bias) + layer = RNN("rnn", n_in=4, n_hidden=3, bias=False) out = layer(X) W_ih = rng.normal(size=(4, 3)).astype(floatX) W_hh = rng.normal(size=(3, 3)).astype(floatX) - b = rng.normal(size=(3,)).astype(floatX) if bias else np.zeros(3, dtype=floatX) layer.cell.W_ih.set_value(W_ih) layer.cell.W_hh.set_value(W_hh) - if bias: - layer.cell.b.set_value(b) X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - assert set(collect_trainable_params(out)) == ( - {layer.cell.W_ih, layer.cell.W_hh, layer.cell.b} - if bias - else {layer.cell.W_ih, layer.cell.W_hh} - ) + assert set(collect_trainable_params(out)) == {layer.cell.W_ih, layer.cell.W_hh} np.testing.assert_allclose( - out.eval({X: X_np}), unrolled(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL + out.eval({X: X_np}), + unrolled(X_np, W_ih, np.zeros(3, dtype=floatX), W_hh, np.tanh), + atol=ATOL, ) -def test_the_recurrent_weight_takes_its_own_initializer(): - """The recurrent draw has a keyword of its own, and using it must not disturb the input projection, - which shares the layer's other two.""" - layer = RNN("rnn", n_in=4, n_hidden=3, recurrent_initializer=ZeroInitializer()) +@pytest.mark.parametrize( + "layer_type, biases", + [(RNN, ["b"]), (GRU, ["b", "c"]), (LSTM, ["b"])], + ids=["rnn", "gru", "lstm"], +) +def test_every_recurrent_layer_forwards_its_initializers_to_the_cell(layer_type, biases): + """Three keyword-only arguments reach the cell through the flat constructor, and one dropped on the + floor, or handed to the wrong parameter, leaves a parameter silently at some other draw. A distinct + constant per keyword says which one reached which parameter. A GRU's two biases share the one + keyword.""" + layer = layer_type( + "layer", + n_in=4, + n_hidden=3, + weight_initializer=constant(value=1.0), + recurrent_initializer=constant(value=2.0), + bias_initializer=constant(value=3.0), + ) + cell = layer.cell - np.testing.assert_array_equal(layer.cell.W_hh.get_value(), np.zeros((3, 3))) - assert np.abs(layer.cell.W_ih.get_value()).max() > 0.0 + assert np.all(cell.W_ih.get_value() == 1.0) + assert np.all(cell.W_hh.get_value() == 2.0) + for bias in biases: + assert np.all(getattr(cell, bias).get_value() == 3.0) -@pytest.mark.parametrize( - "batch_shape", [(), (5,), (2, 5)], ids=["unbatched", "one_axis", "two_axes"] -) +@pytest.mark.parametrize("batch_shape", [(), (2, 5)], ids=["unbatched", "two_axes"]) def test_recurs_over_any_number_of_batch_axes(batch_shape, rng): """Time is the second-to-last axis, as it is for every other layer here. Taking the batch axis to be the leading one instead would give the right answer for a single batch axis and quietly transpose a @@ -231,22 +170,28 @@ def test_recurs_over_any_number_of_batch_axes(batch_shape, rng): np.testing.assert_allclose(result, unrolled(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL) -def test_the_state_takes_the_dtype_the_step_produces(): - """A float32 network fed a float64 sequence. The step promotes, so the state has to promote with it; - a state pinned to floatX leaves scan comparing float32 against the float64 its inner function returns - and refusing the graph. Nothing else here catches it, because every other test runs at one dtype.""" - with pytensor.config.change_flags(floatX="float32"): - layer = RNN("rnn", n_in=4, n_hidden=3) - X = pt.tensor("X", shape=(None, None, 4), dtype="float64") - out = layer(X) +@pytest.mark.parametrize("layer_type", [RNN, GRU, LSTM], ids=["rnn", "gru", "lstm"]) +@pytest.mark.parametrize( + "parameter_dtype, input_dtype", + [("float32", "float64"), ("float64", "float32")], + ids=["wider_input", "wider_parameters"], +) +def test_the_state_takes_the_dtype_the_step_produces(layer_type, parameter_dtype, input_dtype): + """The step promotes the input against the parameters, so the zero state has to promote the same + way. A state pinned to floatX fails the wider input, and one built from the input's dtype alone + fails the wider parameters: either leaves scan comparing a float32 state against the float64 its + inner function returns, and it refuses the graph while building it. Every other test here runs at + one dtype.""" + with pytensor.config.change_flags(floatX=parameter_dtype): + layer = layer_type("layer", n_in=4, n_hidden=3) + out = layer(pt.tensor("X", shape=(None, None, 4), dtype=input_dtype)) - assert out.dtype == "float64" - assert out.eval({X: np.zeros((2, 5, 4), dtype="float64")}).dtype == np.dtype("float64") + assert out.dtype == "float64" class TwoStateCell(RecurrentCell): - """A cell carrying more than one tensor, as an LSTM does. It sums the input into the first state and - counts steps in the second, so both have to survive the round trip to give the right answer.""" + """A cell carrying more than one tensor, as an LSTM does, for checking how a starting state is + matched against several.""" def __init__(self, n_hidden): self.n_hidden = n_hidden @@ -259,41 +204,15 @@ def initial_state(self, X): return zeros, pt.zeros_like(zeros) -def test_a_cell_may_carry_more_than_one_state(rng): - """Scan hands back a list once a cell carries several states rather than the bare variable it returns - for one, and only the first is the output. Taking the wrong one would return the step count.""" - X = pt.tensor("X", shape=(None, None, 3)) - out = Recurrent(TwoStateCell(3), name="two_state")(X) - - X_np = rng.normal(size=(5, 7, 3)).astype(floatX) - - np.testing.assert_allclose(out.eval({X: X_np}), np.cumsum(X_np, axis=-2), atol=ATOL) - - -def test_a_cell_carrying_several_states_can_be_started_from_all_of_them(rng): - """Each state has to reach the step it belongs to. Threading them in the wrong order, or dropping all - but the first, would still run and still return something of the right shape.""" - X = pt.tensor("X", shape=(None, None, 3)) - start = pt.tensor("start", shape=(None, 3)) - out = Recurrent(TwoStateCell(3), name="two_state")(X, [start, pt.zeros_like(start)]) - - X_np = rng.normal(size=(5, 7, 3)).astype(floatX) - start_np = rng.normal(size=(5, 3)).astype(floatX) - - np.testing.assert_allclose( - out.eval({X: X_np, start: start_np}), - np.cumsum(X_np, axis=-2) + start_np[:, None, :], - atol=ATOL, - ) - - def test_a_rejected_state_names_which_one_of_several_is_wrong(): """The message carries a position because a cell may carry many states, and a caller staring at two identically shaped arguments needs to know which one is wrong. A hardcoded index would read as 0 here.""" X = pt.tensor("X", shape=(None, None, 3)) good, bad = pt.matrix("good"), pt.vector("bad") - with pytest.raises(ValueError, match="state at position 1; got a 1-dimensional one"): + with pytest.raises( + ValueError, match="needs a 2-dimensional state at position 1; got a 1-dimensional one" + ): Recurrent(TwoStateCell(3), name="two_state")(X, [good, bad]) @@ -306,13 +225,11 @@ def test_rejects_a_starting_state_the_cell_does_not_carry(): Recurrent(TwoStateCell(3), name="two_state")(X, pt.matrix("only_one")) -def test_the_rnn_is_a_recurrent_over_an_elman_cell(): - """The flat constructor is a convenience over the same two pieces, so a caller who wants the cell on - its own gets exactly what the layer would have built.""" +def test_the_rnn_names_its_cell_parameters_after_the_layer(): + """The flat constructor hands its name to the cell it builds, so the parameters it owns are named + for the layer a caller wrote rather than for a cell they never saw.""" layer = RNN("rnn", n_in=4, n_hidden=3) - assert isinstance(layer, Recurrent) - assert isinstance(layer.cell, ElmanCell) assert [p.name for p in (layer.cell.W_ih, layer.cell.b, layer.cell.W_hh)] == [ "rnn_W_ih", "rnn_b", @@ -320,26 +237,6 @@ def test_the_rnn_is_a_recurrent_over_an_elman_cell(): ] -def test_a_hand_built_cell_scans_the_same_as_the_flat_constructor(rng): - """What the split is for: `Recurrent` takes any cell, and wrapping the same one the flat constructor - builds has to give the same graph back.""" - X = pt.tensor("X", shape=(None, None, 4)) - cell = ElmanCell("cell", n_in=4, n_hidden=3) - wrapped = Recurrent(cell, name="wrapped")(X) - - W_ih = rng.normal(size=(4, 3)).astype(floatX) - b = rng.normal(size=(3,)).astype(floatX) - W_hh = rng.normal(size=(3, 3)).astype(floatX) - cell.W_ih.set_value(W_ih) - cell.b.set_value(b) - cell.W_hh.set_value(W_hh) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - - np.testing.assert_allclose( - wrapped.eval({X: X_np}), unrolled(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL - ) - - def test_rejects_an_input_with_no_time_axis(): layer = RNN("rnn", n_in=4, n_hidden=3) @@ -347,19 +244,21 @@ def test_rejects_an_input_with_no_time_axis(): layer(pt.tensor("X", shape=(4,))) -def test_rejects_an_initial_state_that_does_not_match_the_batch_axes(): - """The state carries one value per batch element, so its rank is fixed by the input's. Scan would - otherwise broadcast a mismatched state into the recurrence and return a silently wrong shape.""" - layer = RNN("rnn", n_in=4, n_hidden=3) +def sigmoid(x): + return 1.0 / (1.0 + np.exp(-x)) - with pytest.raises( - ValueError, match="needs a 2-dimensional state at position 0; got a 1-dimensional one" - ): - layer(pt.tensor("X", shape=(None, None, 4)), pt.tensor("h0", shape=(3,))) +def relu(x): + return np.maximum(x, 0.0) -def sigmoid(x): - return 1.0 / (1.0 + np.exp(-x)) + +class HardSigmoid(Activation): + def __call__(self, x): + return pt.clip(x * 0.2 + 0.5, 0.0, 1.0) + + +def hard_sigmoid(x): + return np.clip(x * 0.2 + 0.5, 0.0, 1.0) def unrolled_gru(X_np, W_ih, b, W_hh, c, phi, gate=sigmoid): @@ -396,14 +295,9 @@ def draw_gru_parameters(layer, rng): return W_ih, b, W_hh, c -@pytest.mark.parametrize( - "activation, reference", - [(Tanh(), np.tanh), (ReLU(), lambda x: np.maximum(x, 0.0))], - ids=["tanh", "relu"], -) -def test_the_gru_matches_a_step_by_step_reference(activation, reference, rng): +def test_the_gru_matches_a_step_by_step_reference(rng): X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3, activation=activation) + layer = GRU("gru", n_in=4, n_hidden=3) out = layer(X) assert out.type.shape == (None, None, 3) @@ -411,7 +305,7 @@ def test_the_gru_matches_a_step_by_step_reference(activation, reference, rng): X_np = rng.normal(size=(5, 7, 4)).astype(floatX) np.testing.assert_allclose( - out.eval({X: X_np}), unrolled_gru(X_np, W_ih, b, W_hh, c, reference), atol=ATOL + out.eval({X: X_np}), unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh), atol=ATOL ) @@ -443,35 +337,11 @@ def test_the_gru_gate_slices_do_not_cross(rng): np.testing.assert_allclose(forgotten, np.tanh(X_np @ W_in), atol=1e-6) -def test_the_gru_candidate_bias_sits_inside_the_reset_gate(rng): - """``c`` is a separate parameter from the candidate's slice of ``b`` only because the reset gate - scales it; folded into ``b`` it would survive a shut gate. Everything else is zeroed, so the whole - output is the bias the gate does or does not let through.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3) - out = layer(X) - - layer.cell.W_ih.set_value(np.zeros((4, 9), dtype=floatX)) - layer.cell.W_hh.set_value(np.zeros((3, 9), dtype=floatX)) - c = rng.normal(size=(3,)).astype(floatX) - layer.cell.c.set_value(c) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - - layer.cell.b.set_value(np.array([20, 20, 20, -20, -20, -20, 0, 0, 0], dtype=floatX)) - np.testing.assert_allclose( - out.eval({X: X_np}), np.broadcast_to(np.tanh(c), (5, 7, 3)), atol=1e-6 - ) - - layer.cell.b.set_value(np.array([-20, -20, -20, -20, -20, -20, 0, 0, 0], dtype=floatX)) - np.testing.assert_allclose(out.eval({X: X_np}), np.zeros((5, 7, 3)), atol=1e-6) - - -@pytest.mark.parametrize("bias", [True, False], ids=["bias", "no_bias"]) -def test_the_gru_biases_are_optional(bias, rng): +def test_the_gru_biases_are_optional(rng): """Dropping the bias drops both parameters as well as both terms; an unused one left behind would hand the optimizer moment state to carry for a weight that never moves.""" X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3, bias=bias) + layer = GRU("gru", n_in=4, n_hidden=3, bias=False) out = layer(X) W_ih = rng.normal(size=(4, 9)).astype(floatX) @@ -480,102 +350,21 @@ def test_the_gru_biases_are_optional(bias, rng): layer.cell.W_hh.set_value(W_hh) b = np.zeros(9, dtype=floatX) c = np.zeros(3, dtype=floatX) - if bias: - b = rng.normal(size=(9,)).astype(floatX) - c = rng.normal(size=(3,)).astype(floatX) - layer.cell.b.set_value(b) - layer.cell.c.set_value(c) X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - assert set(collect_trainable_params(out)) == ( - {layer.cell.W_ih, layer.cell.W_hh, layer.cell.b, layer.cell.c} - if bias - else {layer.cell.W_ih, layer.cell.W_hh} - ) + assert set(collect_trainable_params(out)) == {layer.cell.W_ih, layer.cell.W_hh} np.testing.assert_allclose( out.eval({X: X_np}), unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh), atol=ATOL ) -def test_the_gru_recurrent_weight_is_drawn_orthogonal_by_default(): - """One draw covers all three gates, as in keras, so the check is on the whole wide matrix rather - than on each gate's block.""" - layer = GRU("gru", n_in=16, n_hidden=32) - - W_hh = layer.cell.W_hh.get_value() - assert W_hh.shape == (32, 96) - np.testing.assert_allclose(W_hh @ W_hh.T, np.eye(32), atol=ATOL) - - W_ih = layer.cell.W_ih.get_value() - assert np.abs(W_ih @ W_ih.T - np.eye(16)).max() > 0.1 - # Both fans count the stacked axis, so the spread would be wrong if the draw saw one gate's shape. - assert W_ih.std() == pytest.approx(np.sqrt(2.0 / (16 + 96)), rel=0.1) - - -def test_the_gru_forwards_its_initializers_to_the_cell(): - """Four keyword-only arguments reach the cell through the flat constructor, and one dropped on the - floor leaves a parameter silently at its default draw. Both biases share the one keyword.""" - layer = GRU( - "gru", - n_in=4, - n_hidden=3, - recurrent_initializer=ZeroInitializer(), - bias_initializer=OneInitializer(), - ) - - np.testing.assert_array_equal(layer.cell.W_hh.get_value(), np.zeros((3, 9))) - np.testing.assert_array_equal(layer.cell.b.get_value(), np.ones(9)) - np.testing.assert_array_equal(layer.cell.c.get_value(), np.ones(3)) - assert np.abs(layer.cell.W_ih.get_value()).max() > 0.0 - - -def test_the_gru_trains_end_to_end(rng): - """Gradients survive the round trip through the gates and the training machinery moves them.""" - X = Input("X", shape=(None, 6, 4)) - y = Linear("head", n_in=5, n_out=1)(GRU("gru", n_in=4, n_hidden=5)(X)[..., -1, :]) - model = Model(y).initialize(seed=1) - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) - - X_np = rng.normal(size=(32, 6, 4)).astype(floatX) - y_np = X_np.sum(axis=(1, 2))[:, None].astype(floatX) - - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 - - -def test_the_gru_is_a_recurrent_over_a_gru_cell(rng): - """The flat constructor is a convenience over the cell, and has to build the same graph as writing - the two out by hand.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3) - assert isinstance(layer, Recurrent) - assert isinstance(layer.cell, GRUCell) - - by_hand = Recurrent(GRUCell("gru", n_in=4, n_hidden=3), name="gru") - W_ih, b, W_hh, c = draw_gru_parameters(layer, rng) - by_hand.cell.W_ih.set_value(W_ih) - by_hand.cell.W_hh.set_value(W_hh) - by_hand.cell.b.set_value(b) - by_hand.cell.c.set_value(c) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - - np.testing.assert_allclose(layer(X).eval({X: X_np}), by_hand(X).eval({X: X_np}), atol=ATOL) - - def test_the_gru_gates_take_their_own_activation(rng): - """The gates and the candidate have separate keywords, and setting one must leave the other alone. - A hard sigmoid clipped at the same endpoints is the substitution a reader would actually make, and - it disagrees with the logistic everywhere except the two points where they cross.""" - - class HardSigmoid(Activation): - def __call__(self, x): - return pt.clip(x * 0.2 + 0.5, 0.0, 1.0) - - def hard_sigmoid(x): - return np.clip(x * 0.2 + 0.5, 0.0, 1.0) - + """The gates and the candidate have separate keywords, and each has to reach its own place: one + leaking into the other, or a hardcoded default in either, disagrees with the reference. A hard + sigmoid clipped at the same endpoints is the substitution a reader would actually make, and it + disagrees with the logistic everywhere except the two points where they cross.""" X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3, gate_activation=HardSigmoid()) + layer = GRU("gru", n_in=4, n_hidden=3, activation=ReLU(), gate_activation=HardSigmoid()) out = layer(X) W_ih, b, W_hh, c = draw_gru_parameters(layer, rng) @@ -583,44 +372,28 @@ def hard_sigmoid(x): hard = out.eval({X: X_np}) np.testing.assert_allclose( - hard, unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh, gate=hard_sigmoid), atol=ATOL + hard, unrolled_gru(X_np, W_ih, b, W_hh, c, relu, gate=hard_sigmoid), atol=ATOL ) - # The candidate still runs through tanh, so the default cell is a different function, not a rescaling. - logistic = unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh) + # The same candidate behind logistic gates is a different function, not a rescaling. + logistic = unrolled_gru(X_np, W_ih, b, W_hh, c, relu) assert np.abs(hard - logistic).max() > 0.01 -@pytest.mark.parametrize( - "batch_shape", [(), (5,), (2, 5)], ids=["unbatched", "one_axis", "two_axes"] -) -def test_the_gru_recurs_over_any_number_of_batch_axes(batch_shape, rng): +def test_the_gru_recurs_over_any_number_of_batch_axes(rng): """Every gate is a slice of the last axis, and the candidate's bias broadcasts against whatever - batch axes precede it. A bare sequence has none at all, and a stacked batch has two.""" - X = pt.tensor("X", shape=(*(None for _ in batch_shape), None, 4)) + batch axes precede it. A bare sequence has none at all.""" + X = pt.tensor("X", shape=(None, 4)) layer = GRU("gru", n_in=4, n_hidden=3) out = layer(X) W_ih, b, W_hh, c = draw_gru_parameters(layer, rng) - X_np = rng.normal(size=(*batch_shape, 7, 4)).astype(floatX) + X_np = rng.normal(size=(7, 4)).astype(floatX) np.testing.assert_allclose( out.eval({X: X_np}), unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh), atol=ATOL ) -def test_the_gru_state_takes_the_dtype_the_step_produces(): - """A float32 cell fed a float64 sequence. The cell has to hand the state builder every parameter its - step touches; naming too few, or pinning the state to floatX, leaves scan comparing float32 against - the float64 the step returns and refusing the graph. Every other test here runs at one dtype.""" - with pytensor.config.change_flags(floatX="float32"): - layer = GRU("gru", n_in=4, n_hidden=3) - X = pt.tensor("X", shape=(None, None, 4), dtype="float64") - out = layer(X) - - assert out.dtype == "float64" - assert out.eval({X: np.zeros((2, 5, 4), dtype="float64")}).dtype == np.dtype("float64") - - def unrolled_lstm(X_np, W_ih, b, W_hh, phi, gate=sigmoid): """The gated recurrence written as a python loop, as the reference to check the scan against.""" n_hidden = W_hh.shape[0] @@ -650,16 +423,9 @@ def draw_lstm_parameters(layer, rng): return W_ih, b, W_hh -@pytest.mark.parametrize( - "activation, reference", - [(Tanh(), np.tanh), (ReLU(), lambda x: np.maximum(x, 0.0))], - ids=["tanh", "relu"], -) -def test_the_lstm_matches_a_step_by_step_reference(activation, reference, rng): - """The relu case is what pins ``activation`` being applied twice: once to the candidate and again - to the memory on the way out. At tanh alone, hardcoding either one would still agree.""" +def test_the_lstm_matches_a_step_by_step_reference(rng): X = pt.tensor("X", shape=(None, None, 4)) - layer = LSTM("lstm", n_in=4, n_hidden=3, activation=activation) + layer = LSTM("lstm", n_in=4, n_hidden=3) out = layer(X) assert out.type.shape == (None, None, 3) @@ -667,7 +433,7 @@ def test_the_lstm_matches_a_step_by_step_reference(activation, reference, rng): X_np = rng.normal(size=(5, 7, 4)).astype(floatX) np.testing.assert_allclose( - out.eval({X: X_np}), unrolled_lstm(X_np, W_ih, b, W_hh, reference), atol=ATOL + out.eval({X: X_np}), unrolled_lstm(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL ) @@ -703,44 +469,11 @@ def biases(gate_in, gate_forget, candidate, gate_out): np.testing.assert_allclose(out.eval({X: X_np}), np.zeros((5, 7, 3)), atol=1e-6) -def test_the_lstm_forget_gate_decays_the_memory_over_the_sequence(rng): - """The memory is the state the layer exists for, and it is the one the output never shows directly. - Writing it once and then shutting the input gate leaves the forget gate alone with it: held open the - memory survives every later step, held shut it is gone by the next one.""" - X = pt.tensor("X", shape=(None, None, 4)) - h0 = pt.tensor("h0", shape=(None, 3)) - c0 = pt.tensor("c0", shape=(None, 3)) - layer = LSTM("lstm", n_in=4, n_hidden=3) - out = layer(X, [h0, c0]) - - layer.cell.W_ih.set_value(np.zeros((4, 12), dtype=floatX)) - layer.cell.W_hh.set_value(np.zeros((3, 12), dtype=floatX)) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - h0_np = np.zeros((5, 3), dtype=floatX) - c0_np = rng.normal(size=(5, 3)).astype(floatX) - # Input gate shut, output gate open: h is tanh of whatever the memory still holds. - remembering = np.repeat([-20.0, 20.0, 0.0, 20.0], 3).astype(floatX) - - layer.cell.b.set_value(remembering) - held = out.eval({X: X_np, h0: h0_np, c0: c0_np}) - np.testing.assert_allclose( - held, np.broadcast_to(np.tanh(c0_np)[:, None, :], (5, 7, 3)), atol=1e-5 - ) - - forgetting = remembering.copy() - forgetting[3:6] = -20.0 - layer.cell.b.set_value(forgetting) - np.testing.assert_allclose( - out.eval({X: X_np, h0: h0_np, c0: c0_np}), np.zeros((5, 7, 3)), atol=1e-6 - ) - - -@pytest.mark.parametrize("bias", [True, False], ids=["bias", "no_bias"]) -def test_the_lstm_bias_is_optional(bias, rng): +def test_the_lstm_bias_is_optional(rng): """Dropping the bias drops the parameter as well as the term; an unused one left behind would hand the optimizer moment state to carry for a weight that never moves.""" X = pt.tensor("X", shape=(None, None, 4)) - layer = LSTM("lstm", n_in=4, n_hidden=3, bias=bias) + layer = LSTM("lstm", n_in=4, n_hidden=3, bias=False) out = layer(X) W_ih = rng.normal(size=(4, 12)).astype(floatX) @@ -748,33 +481,20 @@ def test_the_lstm_bias_is_optional(bias, rng): layer.cell.W_ih.set_value(W_ih) layer.cell.W_hh.set_value(W_hh) b = np.zeros(12, dtype=floatX) - if bias: - b = rng.normal(size=(12,)).astype(floatX) - layer.cell.b.set_value(b) X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - assert set(collect_trainable_params(out)) == ( - {layer.cell.W_ih, layer.cell.W_hh, layer.cell.b} - if bias - else {layer.cell.W_ih, layer.cell.W_hh} - ) + assert set(collect_trainable_params(out)) == {layer.cell.W_ih, layer.cell.W_hh} np.testing.assert_allclose( out.eval({X: X_np}), unrolled_lstm(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL ) def test_the_lstm_gates_take_their_own_activation(rng): - """The gates and the candidate have separate keywords, and setting one must leave the other alone.""" - - class HardSigmoid(Activation): - def __call__(self, x): - return pt.clip(x * 0.2 + 0.5, 0.0, 1.0) - - def hard_sigmoid(x): - return np.clip(x * 0.2 + 0.5, 0.0, 1.0) - + """The gates and the candidate have separate keywords, and each has to reach its own place. A relu + ``activation`` also pins it being applied twice: once to the candidate and again to the memory on + the way out. At tanh, hardcoding either one would still agree.""" X = pt.tensor("X", shape=(None, None, 4)) - layer = LSTM("lstm", n_in=4, n_hidden=3, gate_activation=HardSigmoid()) + layer = LSTM("lstm", n_in=4, n_hidden=3, activation=ReLU(), gate_activation=HardSigmoid()) out = layer(X) W_ih, b, W_hh = draw_lstm_parameters(layer, rng) @@ -782,86 +502,27 @@ def hard_sigmoid(x): hard = out.eval({X: X_np}) np.testing.assert_allclose( - hard, unrolled_lstm(X_np, W_ih, b, W_hh, np.tanh, gate=hard_sigmoid), atol=ATOL + hard, unrolled_lstm(X_np, W_ih, b, W_hh, relu, gate=hard_sigmoid), atol=ATOL ) - logistic = unrolled_lstm(X_np, W_ih, b, W_hh, np.tanh) + logistic = unrolled_lstm(X_np, W_ih, b, W_hh, relu) assert np.abs(hard - logistic).max() > 0.01 -def test_the_lstm_forwards_its_initializers_to_the_cell(): - """Four keyword-only arguments reach the cell through the flat constructor, and one dropped on the - floor leaves a parameter silently at its default draw.""" - layer = LSTM( - "lstm", - n_in=4, - n_hidden=3, - recurrent_initializer=ZeroInitializer(), - bias_initializer=OneInitializer(), - ) - - np.testing.assert_array_equal(layer.cell.W_hh.get_value(), np.zeros((3, 12))) - np.testing.assert_array_equal(layer.cell.b.get_value(), np.ones(12)) - assert np.abs(layer.cell.W_ih.get_value()).max() > 0.0 - - -def test_the_lstm_recurrent_weight_is_drawn_orthogonal_by_default(): - """One draw covers all four gates, so the check is on the whole wide matrix rather than per gate.""" - layer = LSTM("lstm", n_in=16, n_hidden=32) - - W_hh = layer.cell.W_hh.get_value() - assert W_hh.shape == (32, 128) - np.testing.assert_allclose(W_hh @ W_hh.T, np.eye(32), atol=ATOL) - - W_ih = layer.cell.W_ih.get_value() - assert np.abs(W_ih @ W_ih.T - np.eye(16)).max() > 0.1 - assert W_ih.std() == pytest.approx(np.sqrt(2.0 / (16 + 128)), rel=0.1) - - -def test_the_lstm_trains_end_to_end(rng): - """Gradients survive the round trip through both carried states and the training machinery moves - them. Nothing else here differentiates a cell whose two states feed each other.""" - X = Input("X", shape=(None, 6, 4)) - y = Linear("head", n_in=5, n_out=1)(LSTM("lstm", n_in=4, n_hidden=5)(X)[..., -1, :]) - model = Model(y).initialize(seed=1) - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) - - X_np = rng.normal(size=(32, 6, 4)).astype(floatX) - y_np = X_np.sum(axis=(1, 2))[:, None].astype(floatX) - - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 - - -@pytest.mark.parametrize( - "batch_shape", [(), (5,), (2, 5)], ids=["unbatched", "one_axis", "two_axes"] -) -def test_the_lstm_recurs_over_any_number_of_batch_axes(batch_shape, rng): +def test_the_lstm_recurs_over_any_number_of_batch_axes(rng): """Both carried states take their batch axes from the input, and the memory has to keep them across the step that combines it with the gates.""" - X = pt.tensor("X", shape=(*(None for _ in batch_shape), None, 4)) + X = pt.tensor("X", shape=(None, 4)) layer = LSTM("lstm", n_in=4, n_hidden=3) out = layer(X) W_ih, b, W_hh = draw_lstm_parameters(layer, rng) - X_np = rng.normal(size=(*batch_shape, 7, 4)).astype(floatX) + X_np = rng.normal(size=(7, 4)).astype(floatX) np.testing.assert_allclose( out.eval({X: X_np}), unrolled_lstm(X_np, W_ih, b, W_hh, np.tanh), atol=ATOL ) -def test_the_lstm_state_takes_the_dtype_the_step_produces(): - """A float32 cell fed a float64 sequence. Both carried states go through the one builder, so a state - pinned to floatX leaves scan comparing float32 against the float64 the step returns.""" - with pytensor.config.change_flags(floatX="float32"): - layer = LSTM("lstm", n_in=4, n_hidden=3) - X = pt.tensor("X", shape=(None, None, 4), dtype="float64") - out = layer(X) - - assert out.dtype == "float64" - assert out.eval({X: np.zeros((2, 5, 4), dtype="float64")}).dtype == np.dtype("float64") - - def test_the_lstm_memory_carries_gradient_across_a_long_sequence(rng): """What the memory is for. With the forget gate open it reaches the last step untouched by any weight, so the gradient back to the starting memory survives fifty steps; with the gate shut the @@ -896,13 +557,9 @@ def test_a_reversed_layer_reads_the_sequence_from_the_end(rng): Anything that merely reordered the output would agree with a forward pass on a palindrome and on nothing else, so the input here is drawn.""" X = pt.tensor("X", shape=(None, None, 4)) - forward = RNN("rnn", n_in=4, n_hidden=3) backward = RNN("rnn", n_in=4, n_hidden=3, reverse=True) - W_ih, b, W_hh = draw_parameters(forward, rng) - backward.cell.W_ih.set_value(W_ih) - backward.cell.b.set_value(b) - backward.cell.W_hh.set_value(W_hh) + W_ih, b, W_hh = draw_parameters(backward, rng) X_np = rng.normal(size=(5, 7, 4)).astype(floatX) read_backward = backward(X).eval({X: X_np}) @@ -910,71 +567,14 @@ def test_a_reversed_layer_reads_the_sequence_from_the_end(rng): # The reference reads the flipped sequence forward, then puts the answers back where they came from. on_flipped = unrolled(X_np[..., ::-1, :], W_ih, b, W_hh, np.tanh) np.testing.assert_allclose(read_backward, on_flipped[..., ::-1, :], atol=ATOL) - assert np.abs(read_backward - forward(X).eval({X: X_np})).max() > 0.1 - - -def test_a_reversed_layer_stays_aligned_with_the_input(rng): - """The output keeps the input's time order, so step t is the step that read X[t] in both - directions. Returning the backward pass last step first -- what keras and flax do by default -- is - the shape that silently misaligns a bidirectional concatenation.""" - X = pt.tensor("X", shape=(None, None, 4)) - backward = RNN("rnn", n_in=4, n_hidden=3, reverse=True) - draw_parameters(backward, rng) - - # A backward pass sees only the last step when it starts, so its first output is a function of that - # step alone. Aligned, that output sits at the end of the sequence. - X_np = rng.normal(size=(1, 6, 4)).astype(floatX) - read_backward = backward(X) - - whole = read_backward.eval({X: X_np}) - last_step_alone = read_backward.eval({X: X_np[:, -1:, :]}) - np.testing.assert_allclose(whole[:, -1, :], last_step_alone[:, 0, :], atol=ATOL) @pytest.mark.parametrize("layer_type", [RNN, GRU, LSTM], ids=["rnn", "gru", "lstm"]) -def test_every_recurrent_layer_takes_a_direction(layer_type, rng): - """``reverse`` lives on the loop, not on the cell, so all three flat constructors have to forward - it. One that dropped it would quietly run forward. Both passes come off the one layer, so they are - the same recurrence read two ways rather than two draws that happen to agree.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = layer_type("layer", n_in=4, n_hidden=3, reverse=True) - assert layer.reverse - - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - backward = layer(X).eval({X: X_np}) - - layer.reverse = False - on_flipped = layer(X).eval({X: X_np[..., ::-1, :]}) - np.testing.assert_allclose(backward, on_flipped[..., ::-1, :], atol=ATOL) - assert np.abs(backward - layer(X).eval({X: X_np})).max() > 0.1 - - -def test_bidirectional_gives_each_step_both_halves_of_the_sequence(rng): - """The point of the wrapper: at every step the forward half has read the prefix and the backward - half the suffix, of the same step. Both halves are checked against the python loop rather than - against another layer, so a concatenation joining mismatched steps cannot agree with a reference - carrying the same misalignment.""" - X = pt.tensor("X", shape=(None, None, 4)) - forward = GRU("fwd", n_in=4, n_hidden=3) - backward = GRU("bwd", n_in=4, n_hidden=5) - out = Bidirectional(forward, backward)(X) - assert out.type.shape == (None, None, 8) - - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - both = out.eval({X: X_np}) - - def parameters(layer): - cell = layer.cell - return cell.W_ih.get_value(), cell.b.get_value(), cell.W_hh.get_value(), cell.c.get_value() - - np.testing.assert_allclose( - both[..., :3], unrolled_gru(X_np, *parameters(forward), np.tanh), atol=ATOL - ) - np.testing.assert_allclose( - both[..., 3:], - unrolled_gru(X_np[..., ::-1, :], *parameters(backward), np.tanh)[..., ::-1, :], - atol=ATOL, - ) +def test_every_recurrent_layer_takes_a_direction(layer_type): + """``reverse`` lives on the loop, not on the cell, so each flat constructor has to forward it. One + that dropped it would quietly run forward. What the loop does with it is the reversed-layer test's + business.""" + assert layer_type("layer", n_in=4, n_hidden=3, reverse=True).reverse def test_bidirectional_owns_the_direction_of_both_layers(rng): @@ -983,8 +583,9 @@ def test_bidirectional_owns_the_direction_of_both_layers(rng): one on its own afterwards is unaffected.""" X = pt.tensor("X", shape=(None, None, 4)) forward = GRU("fwd", n_in=4, n_hidden=3, reverse=True) - backward = GRU("bwd", n_in=4, n_hidden=3) + backward = GRU("bwd", n_in=4, n_hidden=5) both = Bidirectional(forward, backward)(X) + assert both.type.shape == (None, None, 8) X_np = rng.normal(size=(5, 7, 4)).astype(floatX) evaluated = both.eval({X: X_np}) @@ -1012,43 +613,6 @@ def test_bidirectional_rejects_one_layer_used_twice(): Bidirectional(layer, layer) -def test_bidirectional_trains_end_to_end(rng): - """Both halves have to reach the optimizer; a wrapper that dropped one would still train, on half - the parameters.""" - X = Input("X", shape=(None, 6, 4)) - both = Bidirectional(GRU("fwd", n_in=4, n_hidden=5), GRU("bwd", n_in=4, n_hidden=5))(X) - y = Linear("head", n_in=10, n_out=1)(both[..., -1, :]) - model = Model(y).initialize(seed=1) - - assert len(collect_trainable_params(y)) == 10 - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) - X_np = rng.normal(size=(32, 6, 4)).astype(floatX) - y_np = X_np.sum(axis=(1, 2))[:, None].astype(floatX) - - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 - - -def test_a_call_may_override_the_layers_own_direction(rng): - """``Bidirectional`` asks each half for a direction per call rather than reaching in and setting it, - so the override has to work both ways round and leave the layer as it found it.""" - X = pt.tensor("X", shape=(None, None, 4)) - layer = GRU("gru", n_in=4, n_hidden=3) - W_ih, b, W_hh, c = draw_gru_parameters(layer, rng) - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - - read_forward = unrolled_gru(X_np, W_ih, b, W_hh, c, np.tanh) - read_backward = unrolled_gru(X_np[..., ::-1, :], W_ih, b, W_hh, c, np.tanh)[..., ::-1, :] - - np.testing.assert_allclose(layer(X, reverse=True).eval({X: X_np}), read_backward, atol=ATOL) - assert not layer.reverse - - layer.reverse = True - np.testing.assert_allclose(layer(X, reverse=False).eval({X: X_np}), read_forward, atol=ATOL) - np.testing.assert_allclose(layer(X).eval({X: X_np}), read_backward, atol=ATOL) - assert layer.reverse - - def pad_to(sequences, padded_length): """Stack ragged sequences into a rectangle, with the mask that says where each one ends.""" batch = len(sequences) @@ -1060,31 +624,13 @@ def pad_to(sequences, padded_length): return padded, mask -def test_a_mask_makes_padding_leave_a_backward_pass_alone(rng): - """The case that needs the mask most. Read backward, the padding is consumed before any real step, - so it spoils every output position and no indexing afterwards recovers them.""" - X = pt.tensor("X", shape=(None, None, 4)) - mask = pt.tensor("mask", shape=(None, None), dtype=bool) - layer = RNN("rnn", n_in=4, n_hidden=3, reverse=True) - draw_parameters(layer, rng) - - real = rng.normal(size=(3, 4)).astype(floatX) - padded, mask_np = pad_to([real], padded_length=6) - alone = layer(X).eval({X: real[None]}) - - np.testing.assert_allclose( - layer(X, mask=mask).eval({X: padded, mask: mask_np})[:, :3], alone, atol=ATOL - ) - assert not np.allclose(layer(X).eval({X: padded})[:, :3], alone, atol=ATOL) - - def test_a_mask_lets_one_batch_hold_sequences_of_different_lengths(rng): """The case padding exists for: every row a different length, run as one rectangle. Each row has to match what it would have given on its own.""" X = pt.tensor("X", shape=(None, None, 4)) mask = pt.tensor("mask", shape=(None, None), dtype=bool) layer = GRU("gru", n_in=4, n_hidden=3, reverse=True) - draw_gru_parameters(layer, rng) + parameters = draw_gru_parameters(layer, rng) lengths = [2, 5, 3] sequences = [rng.normal(size=(length, 4)).astype(floatX) for length in lengths] @@ -1092,8 +638,8 @@ def test_a_mask_lets_one_batch_hold_sequences_of_different_lengths(rng): together = layer(X, mask=mask).eval({X: padded, mask: mask_np}) for row, (sequence, length) in enumerate(zip(sequences, lengths)): - alone = layer(X).eval({X: sequence[None]}) - np.testing.assert_allclose(together[row, :length], alone[0], atol=ATOL) + alone = unrolled_gru(sequence[::-1], *parameters, np.tanh)[::-1] + np.testing.assert_allclose(together[row, :length], alone, atol=ATOL) def test_a_mask_holds_every_state_a_cell_carries(rng): @@ -1102,14 +648,14 @@ def test_a_mask_holds_every_state_a_cell_carries(rng): X = pt.tensor("X", shape=(None, None, 4)) mask = pt.tensor("mask", shape=(None, None), dtype=bool) layer = LSTM("lstm", n_in=4, n_hidden=3, reverse=True) - draw_lstm_parameters(layer, rng) + parameters = draw_lstm_parameters(layer, rng) real = rng.normal(size=(4, 4)).astype(floatX) padded, mask_np = pad_to([real], padded_length=9) np.testing.assert_allclose( - layer(X, mask=mask).eval({X: padded, mask: mask_np})[:, :4], - layer(X).eval({X: real[None]}), + layer(X, mask=mask).eval({X: padded, mask: mask_np})[0, :4], + unrolled_lstm(real[::-1], *parameters, np.tanh)[::-1], atol=ATOL, ) @@ -1121,25 +667,21 @@ def test_a_mask_freezes_the_final_state_of_a_padded_forward_pass(rng): X = pt.tensor("X", shape=(None, None, 4)) mask = pt.tensor("mask", shape=(None, None), dtype=bool) layer = RNN("rnn", n_in=4, n_hidden=3) - draw_parameters(layer, rng) + W_ih, b, W_hh = draw_parameters(layer, rng) real = rng.normal(size=(3, 4)).astype(floatX) padded, mask_np = pad_to([real], padded_length=7) out = layer(X, mask=mask).eval({X: padded, mask: mask_np}) - alone = layer(X).eval({X: real[None]}) - last_real = alone[0, -1] + last_real = unrolled(real, W_ih, b, W_hh, np.tanh)[-1] np.testing.assert_allclose(out[0, -1], last_real, atol=ATOL) np.testing.assert_allclose(out[0, 3:], np.broadcast_to(last_real, (4, 3)), atol=ATOL) - unmasked = layer(X).eval({X: padded}) - np.testing.assert_allclose(unmasked[:, :3], alone, atol=ATOL) - assert not np.allclose(unmasked[0, -1], last_real, atol=ATOL) - def test_bidirectional_reads_the_mask_in_both_directions(rng): - """The wrapper's backward half is the one that needs the mask, and it has to reach both halves from - the one argument.""" + """The wrapper's backward half is the one that needs the mask to read the real steps right, and the + forward half needs it to hold its final state across the padding. Both read it from the one + argument.""" X = pt.tensor("X", shape=(None, None, 4)) mask = pt.tensor("mask", shape=(None, None), dtype=bool) forward = GRU("fwd", n_in=4, n_hidden=3) @@ -1147,17 +689,20 @@ def test_bidirectional_reads_the_mask_in_both_directions(rng): layer = Bidirectional(forward, backward) # Drawn, not left at their defaults: a zero bias makes the zero state a fixed point, so the # padding would not move the state and the mask would have nothing to undo. - draw_gru_parameters(forward, rng) - draw_gru_parameters(backward, rng) + forward_parameters = draw_gru_parameters(forward, rng) + backward_parameters = draw_gru_parameters(backward, rng) real = rng.normal(size=(3, 4)).astype(floatX) padded, mask_np = pad_to([real], padded_length=8) together = layer(X, mask=mask).eval({X: padded, mask: mask_np}) - alone_forward = forward(X).eval({X: real[None]}) - alone_backward = backward(X, reverse=True).eval({X: real[None]}) - np.testing.assert_allclose(together[:, :3, :3], alone_forward, atol=ATOL) - np.testing.assert_allclose(together[:, :3, 3:], alone_backward, atol=ATOL) + alone_forward = unrolled_gru(real, *forward_parameters, np.tanh) + alone_backward = unrolled_gru(real[::-1], *backward_parameters, np.tanh)[::-1] + np.testing.assert_allclose(together[0, :3, :3], alone_forward, atol=ATOL) + np.testing.assert_allclose(together[0, :3, 3:], alone_backward, atol=ATOL) + # A forward pass is right on the real steps with or without the mask. Only the padding after them + # shows whether the mask reached this half. + np.testing.assert_allclose(together[0, -1, :3], alone_forward[-1], atol=ATOL) def test_rejects_a_mask_that_does_not_match_the_batch_axes(): @@ -1185,31 +730,6 @@ def test_rejects_a_mask_whose_time_axis_is_shorter_than_the_input(rng): out.eval({X: X_np, mask: np.ones((2, 5), dtype=bool)}) -def test_a_mask_keeps_padding_out_of_the_gradient(rng): - """Training on a padded batch has to move the parameters exactly as training on the sequences alone - would. The masked step still computes -- the switch only discards its result -- so a gradient that - leaked through the discarded branch would make the padding length a hyperparameter. Read backward, - because a forward pass reaches the real steps before any padding and would agree either way.""" - X = pt.tensor("X", shape=(None, None, 4)) - mask = pt.tensor("mask", shape=(None, None), dtype=bool) - layer = RNN("rnn", n_in=4, n_hidden=3, reverse=True) - draw_parameters(layer, rng) - parameters = [layer.cell.W_ih, layer.cell.W_hh, layer.cell.b] - - real = rng.normal(size=(3, 4)).astype(floatX) - padded, mask_np = pad_to([real], padded_length=7) - - over_padded = pt.grad(layer(X, mask=mask)[:, :3].sum(), parameters) - over_real = pt.grad(layer(X).sum(), parameters) - - for padded_gradient, real_gradient in zip(over_padded, over_real): - np.testing.assert_allclose( - padded_gradient.eval({X: padded, mask: mask_np}), - real_gradient.eval({X: real[None]}), - atol=ATOL, - ) - - def test_non_finite_padding_leaves_the_gradient_finite(rng): """Padding is only a placeholder, so a batch padded with NaN has to train exactly like one padded with zeros. The masked step's discarded branch still meets the padding on the backward pass, where @@ -1230,26 +750,6 @@ def test_non_finite_padding_leaves_the_gradient_finite(rng): np.testing.assert_array_equal(nan_gradient, zero_gradient) -def test_a_mask_may_skip_a_step_in_the_middle(rng): - """A mask says which steps count, not how many, so it can drop one from the middle -- which a - per-example length cannot express. The recurrence carries on as if that step were not there.""" - X = pt.tensor("X", shape=(None, None, 4)) - mask = pt.tensor("mask", shape=(None, None), dtype=bool) - layer = RNN("rnn", n_in=4, n_hidden=3) - draw_parameters(layer, rng) - - full = rng.normal(size=(5, 4)).astype(floatX) - skipped = 2 - mask_np = np.ones((1, 5), dtype=bool) - mask_np[0, skipped] = False - - with_gap = layer(X, mask=mask).eval({X: full[None], mask: mask_np}) - without = layer(X).eval({X: np.delete(full, skipped, axis=0)[None]}) - - np.testing.assert_allclose(with_gap[:, :skipped], without[:, :skipped], atol=ATOL) - np.testing.assert_allclose(with_gap[:, skipped + 1 :], without[:, skipped:], atol=ATOL) - - class MatrixMemoryCell(RecurrentCell): """A cell whose state carries two feature axes, as a matrix-memory recurrence does. It sums the step's input into every entry, so a step that ran shows up everywhere in the state.""" @@ -1265,23 +765,6 @@ def initial_state(self, X): return (pt.zeros((*X.shape[:-2], self.rows, self.cols), dtype=X.dtype),) -def test_a_state_may_carry_more_than_one_feature_axis(rng): - """Nothing here constrains a carried state to a single feature axis, so time has to land directly - after the input's batch axes rather than second-to-last -- those differ the moment a cell carries a - matrix. The output is the running sum of each step's input, one entry per state cell.""" - X = pt.tensor("X", shape=(None, None, 4)) - out = Recurrent(MatrixMemoryCell(2, 3), name="matrix")(X) - - X_np = rng.normal(size=(5, 7, 4)).astype(floatX) - evaluated = out.eval({X: X_np}) - - assert evaluated.shape == (5, 7, 2, 3) - running = np.cumsum(X_np.sum(axis=-1), axis=-1) - np.testing.assert_allclose( - evaluated, np.broadcast_to(running[:, :, None, None], (5, 7, 2, 3)), atol=ATOL - ) - - def test_a_mask_holds_a_state_of_any_rank(rng): """The mask names batch elements and the state adds however many feature axes the cell wants, so the two are lined up by the state's rank rather than by assuming exactly one feature axis.""" @@ -1291,9 +774,12 @@ def test_a_mask_holds_a_state_of_any_rank(rng): real = rng.normal(size=(3, 4)).astype(floatX) padded, mask_np = pad_to([real], padded_length=6) + evaluated = layer(X, mask=mask).eval({X: padded, mask: mask_np}) + # Every entry holds the running sum of each step's input, frozen once the real steps run out. + running = np.cumsum(real.sum(axis=-1)) + held = np.concatenate([running, np.full(3, running[-1])]) + assert evaluated.shape == (1, 6, 2, 3) np.testing.assert_allclose( - layer(X, mask=mask).eval({X: padded, mask: mask_np})[:, :3], - layer(X).eval({X: real[None]}), - atol=ATOL, + evaluated, np.broadcast_to(held[None, :, None, None], (1, 6, 2, 3)), atol=ATOL ) From e2963e754c52e193ebfb66e85f0410ac897890e4 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:15:48 -0500 Subject: [PATCH 05/10] Check convolution against independent references only --- tests/test_conv.py | 583 ++++++++++----------------------------------- 1 file changed, 123 insertions(+), 460 deletions(-) diff --git a/tests/test_conv.py b/tests/test_conv.py index 54dd937..ff60aff 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -6,25 +6,19 @@ import pytest from pytensor.compile.builders import OpFromGraph -from pytensor.compile.mode import get_default_mode from pytensor.gradient import verify_grad +from pytensor.graph import rewrite_graph +from pytensor.graph.traversal import apply_ancestors from scipy.signal import correlate from pytensor_ml.activations import ReLU from pytensor_ml.layers import BatchNorm, Conv1D, Conv2D, Flatten, Input, Linear, MaxPool2D -from pytensor_ml.layers.conv import ( - Col2Im, - ConvLayer, - ConvLayerGrad, - Im2Col, - _extract_patches, - _scatter_patches, -) +from pytensor_ml.layers.conv import Col2Im, ConvLayer, ConvLayerGrad, Im2Col from pytensor_ml.loss import SquaredError from pytensor_ml.model import Model from pytensor_ml.optim import adam from pytensor_ml.pytensorf import collect_trainable_params -from pytensor_ml.state import OneInitializer, ZeroInitializer, fans +from pytensor_ml.state import OneInitializer, ZeroInitializer floatX = pytensor.config.floatX @@ -38,81 +32,6 @@ def rng(): return np.random.default_rng(sum(map(ord, "pytensor_ml conv"))) -def test_patches_over_one_spatial_axis_are_the_windows_a_kernel_visits(): - """The windows written out by hand. A sequence 0..5 with width 3 and unit stride visits four of - them, and each row of the result is the three consecutive values the kernel would multiply.""" - X = pt.tensor("X", shape=(None, None, 1)) - patches = _extract_patches(X, kernel_size=(3,), stride=(1,), dilation=(1,)) - - X_np = np.arange(6, dtype=floatX).reshape(1, 6, 1) - got = patches.eval({X: X_np}) - - assert got.shape == (1, 4, 3, 1) - expected = np.array([[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5]], dtype=floatX) - np.testing.assert_array_equal(got[0, :, :, 0], expected) - - -def test_patches_step_by_the_stride_and_skip_by_the_dilation(): - """Stride moves the window's start, dilation spreads its taps, and the two are independent. Read - off 0..7: stride 2 starts at 0, 2, 4; dilation 2 takes every other element within a window.""" - X = pt.tensor("X", shape=(None, None, 1)) - X_np = np.arange(8, dtype=floatX).reshape(1, 8, 1) - - strided = _extract_patches(X, kernel_size=(2,), stride=(3,), dilation=(1,)).eval({X: X_np}) - np.testing.assert_array_equal(strided[0, :, :, 0], np.array([[0, 1], [3, 4], [6, 7]])) - - dilated = _extract_patches(X, kernel_size=(3,), stride=(1,), dilation=(2,)).eval({X: X_np}) - np.testing.assert_array_equal( - dilated[0, :, :, 0], np.array([[0, 2, 4], [1, 3, 5], [2, 4, 6], [3, 5, 7]]) - ) - - -def test_patches_over_two_spatial_axes_pair_every_row_window_with_every_column_window(): - """The 2-D case is the 1-D case on each axis and the product across them, which is what the - broadcast between the two index arrays has to produce. A 3x3 image with a 2x2 window visits four - positions, and the corner ones are the top-left and bottom-right 2x2 blocks.""" - X = pt.tensor("X", shape=(None, None, None, 1)) - patches = _extract_patches(X, kernel_size=(2, 2), stride=(1, 1), dilation=(1, 1)) - - X_np = np.arange(9, dtype=floatX).reshape(1, 3, 3, 1) - got = patches.eval({X: X_np}) - - assert got.shape == (1, 2, 2, 2, 2, 1) - np.testing.assert_array_equal(got[0, 0, 0, :, :, 0], np.array([[0, 1], [3, 4]])) - np.testing.assert_array_equal(got[0, 1, 1, :, :, 0], np.array([[4, 5], [7, 8]])) - - -def test_patches_take_a_different_extent_per_spatial_axis(rng): - """Nothing forces the window to be square, and an implementation that broadcast one axis's index - array over both would agree with a square kernel and disagree here.""" - X = pt.tensor("X", shape=(None, None, None, None)) - patches = _extract_patches(X, kernel_size=(2, 3), stride=(1, 1), dilation=(1, 1)) - - X_np = rng.normal(size=(2, 5, 7, 3)).astype(floatX) - got = patches.eval({X: X_np}) - - assert got.shape == (2, 4, 5, 2, 3, 3) - np.testing.assert_allclose(got[1, 2, 3], X_np[1, 2:4, 3:6], atol=1e-12) - - -def test_patches_carry_every_batch_and_channel_axis_through(rng): - """The gather touches the spatial axes alone, so a batch element and a channel are along for the - ride. Checked against numpy slicing at a window chosen off the diagonal.""" - X = pt.tensor("X", shape=(None, None, None, None)) - patches = _extract_patches(X, kernel_size=(3, 3), stride=(2, 2), dilation=(1, 1)) - - X_np = rng.normal(size=(4, 9, 9, 5)).astype(floatX) - got = patches.eval({X: X_np}) - - assert got.shape == (4, 4, 4, 3, 3, 5) - for batch, row, col in ((0, 0, 3), (3, 2, 1)): - np.testing.assert_allclose( - got[batch, row, col], - X_np[batch, 2 * row : 2 * row + 3, 2 * col : 2 * col + 3], - atol=1e-12, - ) - - def correlate_reference(X_np, W_np, stride=1, dilation=1): """A convolution written with scipy, one channel pair at a time, as the independent reference. @@ -149,133 +68,46 @@ def correlate_reference(X_np, W_np, stride=1, dilation=1): return stacked[(slice(None), *(slice(None, None, step) for step in strides), slice(None))] -@pytest.mark.parametrize( - "spatial, kernel_size", - [((9,), (3,)), ((9, 9), (3, 3)), ((7, 9), (2, 4))], - ids=["1d", "2d_square", "2d_rectangular"], -) -def test_the_conv_op_correlates_like_scipy(spatial, kernel_size, rng): - """The load-bearing correctness test. Every channel pair, summed over input channels, against - scipy's correlation rather than against another pytensor graph.""" +def test_the_conv_op_correlates_like_scipy(rng): + """Every channel pair, summed over input channels, against scipy's correlation rather than against + another pytensor graph. The layers check ranks 1 and 2, so the op is checked at rank 3, which is + also the regression: a dispatch that handles some ranks and refuses the rest does not fall back, it + fails to compile.""" in_channels, out_channels = 3, 4 - X = pt.tensor("X", shape=(None, *(None for _ in spatial), in_channels)) - W = pt.tensor("W", shape=(*kernel_size, in_channels, out_channels)) - op = ConvLayer( - kernel_size=kernel_size, stride=(1,) * len(spatial), dilation=(1,) * len(spatial) - ) + X = pt.tensor("X", shape=(None, None, None, None, in_channels)) + W = pt.tensor("W", shape=(2, 2, 2, in_channels, out_channels)) + op = ConvLayer(kernel_size=(2, 2, 2), stride=(1, 1, 1), dilation=(1, 1, 1)) - X_np = rng.normal(size=(2, *spatial, in_channels)).astype(floatX) - W_np = rng.normal(size=(*kernel_size, in_channels, out_channels)).astype(floatX) + X_np = rng.normal(size=(2, 6, 6, 6, in_channels)).astype(floatX) + W_np = rng.normal(size=(2, 2, 2, in_channels, out_channels)).astype(floatX) np.testing.assert_allclose( op(X, W).eval({X: X_np, W: W_np}), correlate_reference(X_np, W_np), atol=ATOL ) -def test_the_conv_op_does_not_flip_its_kernel(rng): - """Correlation, not convolution. An asymmetric kernel is the only thing that tells them apart, and - a flipped implementation would still agree with scipy's ``convolve``, so pin the convention.""" - X = pt.tensor("X", shape=(None, None, 1)) - W = pt.tensor("W", shape=(2, 1, 1)) - op = ConvLayer(kernel_size=(2,), stride=(1,), dilation=(1,)) - - X_np = np.array([[[1.0], [0.0], [0.0]]], dtype=floatX) - W_np = np.array([[[1.0]], [[10.0]]], dtype=floatX) - - # The kernel's first tap lands on the input's first element; flipped, the 10 would land there. - np.testing.assert_allclose(op(X, W).eval({X: X_np, W: W_np})[0, :, 0], [1.0, 0.0], atol=ATOL) - - -def test_the_conv_op_adds_an_optional_bias(rng): - """The bias is a third input rather than a separate Elemwise outside the op, so that a backend - kernel that fuses it has something to fuse.""" - X = pt.tensor("X", shape=(None, None, 2)) - W = pt.tensor("W", shape=(3, 2, 4)) - b = pt.tensor("b", shape=(4,)) - op = ConvLayer(kernel_size=(3,), stride=(1,), dilation=(1,)) - - X_np = rng.normal(size=(2, 8, 2)).astype(floatX) - W_np = rng.normal(size=(3, 2, 4)).astype(floatX) - b_np = rng.normal(size=(4,)).astype(floatX) - - without = op(X, W).eval({X: X_np, W: W_np}) - with_bias = op(X, W, b).eval({X: X_np, W: W_np, b: b_np}) - np.testing.assert_allclose(with_bias, without + b_np, atol=ATOL) - - def test_the_conv_op_takes_a_gradient_matching_finite_differences(rng): - """Overlapping windows make the patch gather's pullback a scatter-add, and pytensor will produce a - gradient for an advanced-indexing expression whether or not the accumulation is right. Checked - against finite differences in float64, where the step size is meaningful.""" - with pytensor.config.change_flags(floatX="float64"): - X = pt.tensor("X", shape=(None, None, None, 2), dtype="float64") - W = pt.tensor("W", shape=(2, 2, 2, 3), dtype="float64") - op = ConvLayer(kernel_size=(2, 2), stride=(1, 1), dilation=(1, 1)) - - X_np = rng.normal(size=(2, 5, 5, 2)) - W_np = rng.normal(size=(2, 2, 2, 3)) - cost = (op(X, W) ** 2).sum() - - for wrt, value in ((X, X_np), (W, W_np)): - analytic = pt.grad(cost, wrt).eval({X: X_np, W: W_np}) - numeric = np.zeros_like(value) - step = 1e-6 - flat = numeric.reshape(-1) - for index in range(flat.size): - nudged = value.copy().reshape(-1) - nudged[index] += step - up = cost.eval({X: X_np, W: W_np} | {wrt: nudged.reshape(value.shape)}) - nudged[index] -= 2 * step - down = cost.eval({X: X_np, W: W_np} | {wrt: nudged.reshape(value.shape)}) - flat[index] = (up - down) / (2 * step) - np.testing.assert_allclose(analytic, numeric, rtol=1e-5, atol=1e-6) - - -def test_conv1d_correlates_like_scipy_and_adds_its_bias(rng): - """The layer is padding, parameters and a call. Checked against the reference rather than against - the op, so a layer that padded the wrong axis or transposed its kernel would show up.""" - X = pt.tensor("X", shape=(None, None, 3)) - layer = Conv1D("conv", in_channels=3, out_channels=5, kernel_size=4) - out = layer(X) - - W_np = rng.normal(size=(4, 3, 5)).astype(floatX) - b_np = rng.normal(size=(5,)).astype(floatX) - layer.W.set_value(W_np) - layer.b.set_value(b_np) - X_np = rng.normal(size=(2, 11, 3)).astype(floatX) - - np.testing.assert_allclose( - out.eval({X: X_np}), correlate_reference(X_np, W_np) + b_np, atol=ATOL - ) - - -@pytest.mark.parametrize( - "kernel_size, padding, expected", - [(3, "valid", 8), (3, "same", 10), (4, "same", 10), (3, 2, 12)], - ids=["valid", "same_odd", "same_even", "explicit"], -) -def test_conv1d_pads_to_the_length_it_promises(kernel_size, padding, expected, rng): - """`same` is the one that has to hold at even kernel sizes, where the padding cannot be symmetric - and the output length is the thing the name is a claim about.""" - X = pt.tensor("X", shape=(None, None, 2)) - layer = Conv1D("conv", in_channels=2, out_channels=3, kernel_size=kernel_size, padding=padding) + """Overlapping windows make the gather's pullback a scatter-add, and the bias's is a sum over + everything but the channel axis. Checked against finite differences for all three inputs, at a + stride and a dilation that differ per axis, with spatial extents known only at runtime.""" + op = ConvLayer(kernel_size=(2, 2), stride=(2, 1), dilation=(1, 2)) + inputs = [ + rng.normal(size=(2, 7, 7, 2)), + rng.normal(size=(2, 2, 2, 3)), + rng.normal(size=(3,)), + ] - got = layer(X).eval({X: rng.normal(size=(2, 10, 2)).astype(floatX)}) - assert got.shape == (2, expected, 3) + with pytensor.config.change_flags(floatX="float64"): + verify_grad(op, inputs, rng=np.random.default_rng(0)) def test_conv1d_pads_with_the_mode_it_is_given(): """A padding mode reaches `pt.pad`, so an edge-padded input repeats its boundary rather than - fading to zero. With a summing kernel over a constant input, zero padding shows at the edges and - edge padding does not.""" + fading to zero. With a summing kernel over a constant input, zero padding would leave each edge + window short by one, and edge padding leaves none short.""" X = pt.tensor("X", shape=(None, None, 1)) X_np = np.ones((1, 6, 1), dtype=floatX) - ones = np.ones((3, 1, 1), dtype=floatX) - zero_padded = Conv1D( - "c", in_channels=1, out_channels=1, kernel_size=3, padding="same", bias=False - ) - zero_padded.W.set_value(ones) edge_padded = Conv1D( "c", in_channels=1, @@ -285,47 +117,28 @@ def test_conv1d_pads_with_the_mode_it_is_given(): padding_mode="edge", bias=False, ) - edge_padded.W.set_value(ones) + edge_padded.W.set_value(np.ones((3, 1, 1), dtype=floatX)) - np.testing.assert_allclose( - zero_padded(X).eval({X: X_np})[0, :, 0], [2, 3, 3, 3, 3, 2], atol=ATOL - ) np.testing.assert_allclose( edge_padded(X).eval({X: X_np})[0, :, 0], [3, 3, 3, 3, 3, 3], atol=ATOL ) -@pytest.mark.parametrize("bias", [True, False], ids=["bias", "no_bias"]) -def test_conv1d_bias_is_optional(bias, rng): +def test_conv1d_bias_is_optional(): """Dropping the bias drops the parameter as well as the term; an unused one would hand the optimizer moment state to carry for a weight that never moves.""" - X = pt.tensor("X", shape=(None, None, 2)) - layer = Conv1D("conv", in_channels=2, out_channels=3, kernel_size=2, bias=bias) - out = layer(X) + layer = Conv1D("conv", in_channels=2, out_channels=3, kernel_size=2, bias=False) + out = layer(pt.tensor("X", shape=(None, None, 2))) - W_np = rng.normal(size=(2, 2, 3)).astype(floatX) - layer.W.set_value(W_np) - shift = np.zeros(3, dtype=floatX) - if bias: - shift = rng.normal(size=(3,)).astype(floatX) - layer.b.set_value(shift) - X_np = rng.normal(size=(2, 7, 2)).astype(floatX) + assert set(collect_trainable_params(out)) == {layer.W} - assert set(collect_trainable_params(out)) == ({layer.W, layer.b} if bias else {layer.W}) - np.testing.assert_allclose( - out.eval({X: X_np}), correlate_reference(X_np, W_np) + shift, atol=ATOL - ) - -def test_conv1d_draws_its_kernel_with_the_receptive_field_in_the_fans(): - """A conv kernel is stored input-channels-first so `fans` reads it correctly: fan_in is - `in_channels * kernel_size`, not two kernel extents. A kernel-first layout would give a draw scaled - for the wrong fans and nothing else here would notice.""" +def test_conv1d_stores_its_kernel_in_the_layout_fans_reads(): + """`fans` reads the trailing two axes of a kernel as its channels and the rest as the receptive + field, so a kernel stored in any other layout would be drawn at a scale set by the wrong fans.""" layer = Conv1D("conv", in_channels=8, out_channels=16, kernel_size=5) assert layer.W.get_value().shape == (5, 8, 16) - assert fans(layer.W.get_value().shape) == (8 * 5, 16 * 5) - assert layer.W.get_value().std() == pytest.approx(np.sqrt(2.0 / (40 + 80)), rel=0.15) def test_conv1d_forwards_its_initializers_to_its_parameters(): @@ -353,25 +166,6 @@ def test_conv1d_rejects_an_input_of_the_wrong_rank(): layer(pt.tensor("X", shape=(None, 2))) -def test_conv1d_trains_end_to_end(rng): - """Gradients survive the gather, the matmul and the training machinery, and the parameters move. - - The head pools over time rather than taking the last position, because a `valid` convolution's last - output sees only the final window -- a target summing the whole sequence would be unlearnable by - construction, and the test would be measuring nothing.""" - X = Input("X", shape=(None, 12, 2)) - features = Conv1D("conv", in_channels=2, out_channels=4, kernel_size=3, padding="same")(X) - y = Linear("head", n_in=4, n_out=1)(features.sum(axis=1)) - model = Model(y).initialize(seed=1) - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) - - X_np = rng.normal(size=(32, 12, 2)).astype(floatX) - y_np = X_np.sum(axis=(1, 2))[:, None].astype(floatX) - - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 - - def test_conv1d_rejects_a_kernel_wider_than_the_input(): """A window that does not fit yields no windows at all, and every downstream shape then carries a zero axis, so the graph computes an empty answer instead of failing. Caught wherever the input's @@ -386,11 +180,9 @@ def test_conv1d_rejects_a_kernel_wider_than_the_input(): with pytest.raises(ValueError, match="at least 9 elements"): dilated(pt.tensor("X", shape=(None, 8, 1))) - # Padding is counted, so the same kernel fits once there is enough of it. Checked by evaluating: - # the gather does not propagate a static spatial size, so the built graph's type says None here. + # Padding is counted, so the same kernel fits once there is enough of it. padded = Conv1D("conv", in_channels=1, out_channels=1, kernel_size=10, padding=3) - X = pt.tensor("X", shape=(None, 5, 1)) - assert padded(X).eval({X: np.zeros((2, 5, 1), dtype=floatX)}).shape == (2, 2, 1) + assert padded(pt.tensor("X", shape=(None, 5, 1))).type.shape == (None, 2, 1) def test_conv1d_rejects_negative_padding(): @@ -400,29 +192,25 @@ def test_conv1d_rejects_negative_padding(): Conv1D("conv", in_channels=1, out_channels=1, kernel_size=3, padding=-1) -@pytest.mark.parametrize("stride", [1, 2, 3], ids=["stride1", "stride2", "stride3"]) -@pytest.mark.parametrize("length", [10, 11], ids=["even", "odd"]) -def test_conv1d_same_padding_holds_at_any_stride(stride, length, rng): +@pytest.mark.parametrize( + "length, stride", [(10, 3), (11, 2), (11, 3)], ids=["10_stride3", "11_stride2", "11_stride3"] +) +def test_conv1d_same_padding_holds_at_any_stride(length, stride): """`same` is a claim about the output length, and torch and keras both define it as - ceil(input / stride) rather than only holding at unit stride.""" - X = pt.tensor("X", shape=(None, None, 1)) + ceil(input / stride) rather than only holding at unit stride. Only a stride that does not divide + the length separates the ceiling from the floor.""" layer = Conv1D( "conv", in_channels=1, out_channels=1, kernel_size=3, stride=stride, padding="same" ) - got = layer(X).eval({X: rng.normal(size=(2, length, 1)).astype(floatX)}) - assert got.shape == (2, -(-length // stride), 1) + out = layer(pt.tensor("X", shape=(None, length, 1))) + assert out.type.shape == (None, -(-length // stride), 1) -@pytest.mark.parametrize( - "stride, dilation", - [(2, 1), (3, 1), (1, 2), (1, 3), (2, 2)], - ids=["stride2", "stride3", "dilation2", "dilation3", "both"], -) -def test_conv1d_strides_and_dilates_like_scipy(stride, dilation, rng): - """Stride and dilation are checked at the helper against hand-written windows, but nothing until - here checks that they survive the trip from the layer's constructor through the op to the gather. - A layer that swapped the two, or dropped either, would pass every other test in this file.""" +def test_conv1d_strides_and_dilates_like_scipy(rng): + """Stride and dilation have to survive the trip from the layer's constructor through the op to the + gather. They differ here, so a layer that swapped the two, or dropped either, would disagree.""" + stride, dilation = 2, 3 X = pt.tensor("X", shape=(None, None, 3)) layer = Conv1D( "conv", in_channels=3, out_channels=4, kernel_size=3, stride=stride, dilation=dilation @@ -453,97 +241,28 @@ def test_conv1d_rejects_a_per_axis_argument_of_the_wrong_length(): Conv1D("conv", in_channels=1, out_channels=1, kernel_size=3, stride=(1, 1)) -def test_a_kernel_of_one_is_a_linear_layer_applied_per_step(rng): - """A one-tap convolution is a matrix applied at every position, so it has to agree with `Linear` on - the same weights. That pins the two conventions together: the kernel's trailing axes are the same - ``(in, out)`` a weight matrix is, so copying one into the other needs a new axis and nothing else.""" - X = pt.tensor("X", shape=(None, None, 4)) - linear = Linear("dense", n_in=4, n_out=6) - conv = Conv1D("conv", in_channels=4, out_channels=6, kernel_size=1) - - W_np = rng.normal(size=(4, 6)).astype(floatX) - b_np = rng.normal(size=(6,)).astype(floatX) - linear.W.set_value(W_np) - linear.b.set_value(b_np) - conv.W.set_value(W_np[None]) - conv.b.set_value(b_np) - - assert fans(conv.W.get_value().shape) == fans(linear.W.get_value().shape) - - X_np = rng.normal(size=(3, 7, 4)).astype(floatX) - np.testing.assert_allclose(conv(X).eval({X: X_np}), linear(X).eval({X: X_np}), atol=ATOL) - - @pytest.mark.parametrize( "spatial, kernel_size, stride, dilation", - [ - ((11,), (3,), (1,), (1,)), - ((11,), (3,), (2,), (2,)), - ((7, 9), (2, 3), (2, 1), (1, 2)), - ((6, 6, 6), (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ], - ids=["1d", "1d_strided_dilated", "2d_asymmetric", "3d"], + [((12,), (3,), (5,), (1,)), ((6, 6, 6), (2, 2, 2), (1, 1, 1), (1, 1, 1))], + ids=["1d_untouched_tail", "3d"], ) -def test_im2col_matches_the_reference_gather_at_every_rank( - spatial, kernel_size, stride, dilation, rng -): - """The op stands in for the advanced-indexing gather, so it has to agree with it exactly, at every - rank and not only the ranks a layer happens to build. The 3-D case is the regression: a dispatch - that handles some ranks and refuses the rest does not fall back, it fails to compile.""" +def test_col2im_is_the_adjoint_of_im2col(spatial, kernel_size, stride, dilation, rng): + """A scatter-add is the adjoint of the gather it reverses, so `` == `` + for any ``X`` and ``P``. The untouched-tail case is the one a scatter can get wrong on its own + terms: a stride that overshoots leaves positions no window reaches, and those have to stay zero + rather than pick up a neighbor.""" X = pt.tensor("X", shape=(2, *spatial, 3)) - X_np = rng.normal(size=(2, *spatial, 3)).astype(floatX) - - reference = _extract_patches(X, kernel_size, stride, dilation).eval({X: X_np}) - got = pytensor.function([X], Im2Col(kernel_size, stride, dilation)(X))(X_np) - - assert got.shape == reference.shape - np.testing.assert_array_equal(got, reference) - - -def test_im2col_infers_the_shape_it_produces(rng): - """`infer_shape` is what lets the rest of the graph reason about the gather without running it, so - a formula that drifts from `perform` would mis-shape everything downstream and only fail later.""" - X = pt.tensor("X", shape=(None, None, 3)) - patches = Im2Col((3,), (2,), (2,))(X) - X_np = rng.normal(size=(2, 13, 3)).astype(floatX) - - inferred = pytensor.function([X], patches.shape)(X_np) - np.testing.assert_array_equal(inferred, pytensor.function([X], patches)(X_np).shape) - + gathered = Im2Col(kernel_size, stride, dilation)(X) + patches = pt.tensor("patches", shape=gathered.type.shape) + scattered = Col2Im(kernel_size, stride, dilation)(patches, *spatial) + + X_np = rng.normal(size=X.type.shape).astype(floatX) + patches_np = rng.normal(size=patches.type.shape).astype(floatX) + inner_products = pytensor.function( + [X, patches], [(gathered * patches).sum(), (X * scattered).sum()] + ) -@pytest.mark.parametrize( - "spatial, kernel_size, stride, dilation", - [ - ((11,), (3,), (1,), (1,)), - ((11,), (3,), (2,), (2,)), - ((12,), (3,), (5,), (1,)), - ((7, 9), (2, 3), (2, 1), (1, 2)), - ((6, 6, 6), (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ], - ids=["1d", "1d_strided_dilated", "1d_untouched_tail", "2d_asymmetric", "3d"], -) -def test_col2im_matches_the_reference_scatter_at_every_rank( - spatial, kernel_size, stride, dilation, rng -): - """The op stands in for the `inc_subtensor` scatter, so it has to agree with it exactly. The - untouched-tail case is the one a scatter can get wrong on its own terms: a stride that overshoots - leaves positions no window reaches, and those have to stay zero rather than pick up a neighbor.""" - X = pt.tensor("X", shape=(2, *spatial, 3)) - X_np = rng.normal(size=(2, *spatial, 3)).astype(floatX) - patches_np = pytensor.function([X], Im2Col(kernel_size, stride, dilation)(X))(X_np) - - cotangent = pt.tensor("cotangent", shape=patches_np.shape) - reference = pytensor.function( - [cotangent, X], - _scatter_patches(cotangent, X, kernel_size, stride, dilation), - on_unused_input="ignore", - )(patches_np, X_np) - got = pytensor.function( - [cotangent], Col2Im(kernel_size, stride, dilation)(cotangent, *spatial) - )(patches_np) - - assert got.shape == reference.shape - np.testing.assert_allclose(got, reference, atol=ATOL) + np.testing.assert_allclose(*inner_products(X_np, patches_np), rtol=ATOL) def test_col2im_keeps_a_spatial_extent_it_is_given_statically(): @@ -572,21 +291,48 @@ def test_col2im_gathers_the_cotangent_it_scattered(rng): ) -@pytest.mark.parametrize("view", ["transposed", "reversed"]) -def test_im2col_accepts_an_input_it_cannot_assume_is_contiguous(view, rng): +def test_im2col_accepts_an_input_it_cannot_assume_is_contiguous(rng): """`TensorType` carries no layout, so a kernel is typed against any layout whatever the data turns out to be. One that quietly needs a contiguous buffer does not fall back to `perform` when it does not get one -- it fails to compile, for every caller.""" X_np = rng.normal(size=(2, 3, 11)).astype(floatX) X = pt.tensor("X", shape=(2, 3, 11)) - source = X.transpose(0, 2, 1) if view == "transposed" else X[:, :, ::-1].transpose(0, 2, 1) + contiguous = pt.tensor("contiguous", shape=(2, 11, 3)) + im2col = Im2Col((3,), (1,), (1,)) - got = pytensor.function([X], Im2Col((3,), (1,), (1,))(source))(X_np) - reference = _extract_patches(source, (3,), (1,), (1,)).eval({X: X_np}) + got = pytensor.function([X], im2col(X[:, :, ::-1].transpose(0, 2, 1)))(X_np) + reference = pytensor.function([contiguous], im2col(contiguous))( + np.ascontiguousarray(X_np[:, :, ::-1].transpose(0, 2, 1)) + ) np.testing.assert_allclose(got, reference, atol=ATOL) +def test_the_gather_and_scatter_run_the_same_in_python_as_compiled(rng): + """`perform` is what runs wherever no backend dispatches the op, and the compiled default never + reaches it, so the python path is checked against the compiled one. Kernel, stride and dilation + each differ per axis, so a `perform` that swapped any of them across axes would disagree.""" + geometry = ((2, 3), (2, 1), (1, 2)) + X = pt.tensor("X", shape=(None, None, None, 3)) + patches = pt.tensor("patches", shape=(None, None, None, 2, 3, 3)) + height = pt.scalar("height", dtype="int64") + width = pt.scalar("width", dtype="int64") + inputs = [X, patches, height, width] + outputs = [Im2Col(*geometry)(X), Col2Im(*geometry)(patches, height, width)] + + values = ( + rng.normal(size=(2, 7, 8, 3)).astype(floatX), + rng.normal(size=(2, 3, 4, 2, 3, 3)).astype(floatX), + 7, + 8, + ) + in_python = pytensor.function(inputs, outputs, mode="FAST_COMPILE")(*values) + compiled = pytensor.function(inputs, outputs)(*values) + + for python_value, compiled_value in zip(in_python, compiled): + np.testing.assert_allclose(python_value, compiled_value, atol=ATOL) + + def test_col2im_needs_one_extent_per_spatial_axis(): """The extents are positional, so a caller passing the wrong number of them would otherwise build a node whose rank silently disagrees with the kernel's.""" @@ -596,50 +342,11 @@ def test_col2im_needs_one_extent_per_spatial_axis(): Col2Im((3, 3), (1, 1), (1, 1))(patches, 11) -def test_the_gather_backward_runs_when_the_spatial_extent_is_only_known_at_runtime(rng): - """`Im2Col.pullback` hands the input's spatial extents to the scatter as ordinary inputs, so an - extent nobody knows until the graph runs has to arrive as a value rather than as a constant folded - into the node.""" - X_np = rng.normal(size=(2, 11, 3)).astype(floatX) - seed_np = rng.normal(size=(2, 9, 3, 3)).astype(floatX) - seed = pt.tensor("seed", shape=(2, 9, 3, 3)) - - gradients = [] - for spatial in (None, 11): - X = pt.tensor("X", shape=(2, spatial, 3)) - patches = Im2Col((3,), (1,), (1,))(X) - pulled_back = pt.grad(cost=None, wrt=X, known_grads={patches: seed}) - gradients.append(pytensor.function([X, seed], pulled_back)(X_np, seed_np)) - - np.testing.assert_allclose(*gradients, atol=ATOL) - - -def test_conv2d_correlates_like_scipy_and_adds_its_bias(rng): - """The layer at rank 2, against the reference rather than against the op. A rectangular kernel is - what separates the two spatial axes: a layer that transposed them would still pass at rank 1 and - with a square kernel.""" - X = pt.tensor("X", shape=(None, None, None, 3)) - layer = Conv2D("conv", in_channels=3, out_channels=5, kernel_size=(2, 4)) - - W_np = rng.normal(size=(2, 4, 3, 5)).astype(floatX) - b_np = rng.normal(size=(5,)).astype(floatX) - layer.W.set_value(W_np) - layer.b.set_value(b_np) - X_np = rng.normal(size=(2, 9, 11, 3)).astype(floatX) - - np.testing.assert_allclose( - layer(X).eval({X: X_np}), correlate_reference(X_np, W_np) + b_np, atol=ATOL - ) - - -@pytest.mark.parametrize( - "stride, dilation", - [((2, 1), (1, 1)), ((1, 3), (1, 1)), ((1, 1), (2, 1)), ((1, 1), (1, 3)), ((3, 2), (1, 2))], - ids=["stride_h", "stride_w", "dilate_h", "dilate_w", "both_asymmetric"], -) -def test_conv2d_strides_and_dilates_per_axis_like_scipy(stride, dilation, rng): +def test_conv2d_strides_and_dilates_per_axis_like_scipy(rng): """Each axis carries its own stride and dilation, and only an asymmetric setting can catch the two - being swapped or one of them broadcast over both axes.""" + being swapped or one of them broadcast over both axes. The rectangular kernel separates the two + spatial axes, so a layer that transposed them would disagree too.""" + stride, dilation = (3, 2), (1, 2) X = pt.tensor("X", shape=(None, None, None, 3)) layer = Conv2D( "conv", @@ -662,57 +369,22 @@ def test_conv2d_strides_and_dilates_per_axis_like_scipy(stride, dilation, rng): ) -@pytest.mark.parametrize( - "kernel_size, spatial, expected", - [ - ((3, 3), (8, 10), (8, 10)), - ((2, 4), (8, 10), (8, 10)), - ((2, 5), (7, 7), (7, 7)), - ], - ids=["both_odd", "both_even", "mixed_parity"], -) -def test_conv2d_same_padding_holds_each_axis_independently(kernel_size, spatial, expected, rng): - """`same` has no symmetric answer at an even extent, and rank 2 is where one axis can be even while - the other is odd -- the case that catches a padding amount computed once and reused.""" - X = pt.tensor("X", shape=(None, None, None, 3)) - layer = Conv2D("conv", in_channels=3, out_channels=2, kernel_size=kernel_size, padding="same") - X_np = rng.normal(size=(2, *spatial, 3)).astype(floatX) - - assert layer(X).eval({X: X_np}).shape == (2, *expected, 2) +def test_conv2d_pads_each_axis_by_its_own_amount(): + """An explicit padding takes one amount per spatial axis, applied to both sides of that axis, so + a layer that broadcast the first amount or swapped the two would change the output's extents.""" + layer = Conv2D("conv", in_channels=2, out_channels=4, kernel_size=3, padding=(1, 2)) + assert layer(pt.tensor("X", shape=(None, 5, 5, 2))).type.shape == (None, 5, 7, 4) -def test_conv1d_matches_a_conv2d_whose_second_axis_is_degenerate(rng): - """Both layers are `_ConvNd` with a different `n_spatial`, and this is what pins that: a 2-D - convolution over a width-1 image with a width-1 kernel is a 1-D convolution, so the two have to - agree to the last bit rather than merely have matching shapes.""" - W_np = rng.normal(size=(3, 4, 5)).astype(floatX) - b_np = rng.normal(size=(5,)).astype(floatX) - X_np = rng.normal(size=(2, 12, 4)).astype(floatX) - one_d = Conv1D("one_d", in_channels=4, out_channels=5, kernel_size=3, stride=2) - one_d.W.set_value(W_np) - one_d.b.set_value(b_np) +def test_conv2d_rejects_a_padding_it_cannot_read(): + """A misspelled mode and a padding sized for another rank are the natural mistakes, and either + would otherwise surface from `pt.pad` far from the constructor.""" + with pytest.raises(ValueError, match="padding must be 'valid', 'same', or an explicit number"): + Conv2D("conv", in_channels=1, out_channels=1, kernel_size=3, padding="full") - two_d = Conv2D("two_d", in_channels=4, out_channels=5, kernel_size=(3, 1), stride=(2, 1)) - two_d.W.set_value(W_np[:, None]) - two_d.b.set_value(b_np) - - X1 = pt.tensor("X1", shape=(None, None, 4)) - X2 = pt.tensor("X2", shape=(None, None, None, 4)) - np.testing.assert_allclose( - one_d(X1).eval({X1: X_np}), - two_d(X2).eval({X2: X_np[:, :, None]})[:, :, 0], - atol=ATOL, - ) - - -def test_conv2d_draws_its_kernel_with_both_extents_in_the_fans(): - """`fans` takes the two trailing dimensions as the features and everything before them as the - receptive field, so at rank 2 the field is the product of both extents.""" - layer = Conv2D("conv", in_channels=8, out_channels=16, kernel_size=(3, 5)) - - assert layer.W.get_value().shape == (3, 5, 8, 16) - assert fans(layer.W.get_value().shape) == (8 * 15, 16 * 15) + with pytest.raises(ValueError, match="needs one padding amount per axis, but got 3"): + Conv2D("conv", in_channels=1, out_channels=1, kernel_size=3, padding=(1, 1, 1)) def test_conv2d_same_padding_puts_the_odd_element_after_on_each_axis(): @@ -777,33 +449,26 @@ def test_a_convolutional_network_trains_end_to_end(rng): [("dX", (True, False)), ("dW", (False, True)), ("both", (True, True))], ids=["input_only", "kernel_only", "both"], ) -def test_the_pullback_computes_only_the_gradients_something_reads(wanted, expected, rng): +def test_the_pullback_computes_only_the_gradients_something_reads(wanted, expected): """`ConvLayer.pullback` asks for both gradients because only the graph knows which are wanted, and only once it is built. A rewrite then drops whichever has no clients -- the input gradient for the - first convolution of a network, the kernel gradient for a transposed one -- and the numbers it - leaves behind have to be the ones the un-rewritten graph would have produced.""" + first convolution of a network, the kernel gradient for a transposed one. That the gradient it + keeps is unchanged is checked against the op directly, below.""" X = pt.tensor("X", shape=(4, 24, 3)) layer = Conv1D("conv", in_channels=3, out_channels=5, kernel_size=3) - layer.W.set_value(rng.normal(size=(3, 3, 5)).astype(floatX)) - layer.b.set_value(rng.normal(size=(5,)).astype(floatX)) cost = (layer(X) ** 2).sum() targets = { "dX": [pt.grad(cost, X)], "dW": [pt.grad(cost, layer.W)], "both": pt.grad(cost, [X, layer.W]), }[wanted] - X_np = rng.normal(size=(4, 24, 3)).astype(floatX) - fn = pytensor.function([X], targets) + rewritten = rewrite_graph(targets, include=("canonicalize", "specialize")) grad_op = next( - node.op for node in fn.maker.fgraph.apply_nodes if isinstance(node.op, ConvLayerGrad) + node.op for node in apply_ancestors(rewritten) if isinstance(node.op, ConvLayerGrad) ) assert (grad_op.compute_dX, grad_op.compute_dW) == expected - unrewritten = get_default_mode().excluding("drop_unused_input_grad", "drop_unused_kernel_grad") - for dropped, kept in zip(pytensor.function([X], targets, mode=unrewritten)(X_np), fn(X_np)): - np.testing.assert_allclose(kept, dropped, atol=ATOL) - @pytest.mark.parametrize("dropped", ["dX", "dW"], ids=["without_dX", "without_dW"]) def test_dropping_one_gradient_leaves_the_other_unchanged(dropped, rng): @@ -869,14 +534,12 @@ def test_differentiating_a_conv_gradient_leaves_only_dispatchable_ops(flags): [{"compute_dW": False}, {"compute_dX": False}, {}], ids=["dX_only", "dW_only", "both"], ) -@pytest.mark.parametrize( - "stride, dilation", [(1, 1), (2, 1), (1, 2)], ids=["plain", "strided", "dilated"] -) -def test_the_pullback_of_the_pullback_matches_finite_differences(flags, stride, dilation, rng): +def test_the_pullback_of_the_pullback_matches_finite_differences(flags, rng): """The closed forms the pullback uses are adjoint identities rather than a differentiated graph, so they are checked numerically. Every input is perturbed, including the one each output - ignores.""" - geometry = ((3, 3), (stride, stride), (dilation, dilation)) + ignores. The pullback has no branch on geometry, so one stride and one dilation, each differing + per axis, stand in for all of them.""" + geometry = ((3, 3), (2, 1), (1, 2)) X_np = rng.normal(size=(2, 7, 7, 3)).astype(floatX) W_np = rng.normal(size=(3, 3, 3, 4)).astype(floatX) cotangent_shape = ConvLayer(*geometry)( From 8f1a5a466197c2be61720be8d631c74c0126d21a Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:15:56 -0500 Subject: [PATCH 06/10] Remove the test-only patch gather and scatter --- pytensor_ml/layers/conv.py | 69 -------------------------------------- 1 file changed, 69 deletions(-) diff --git a/pytensor_ml/layers/conv.py b/pytensor_ml/layers/conv.py index fc26ffc..6dbf02d 100644 --- a/pytensor_ml/layers/conv.py +++ b/pytensor_ml/layers/conv.py @@ -58,45 +58,6 @@ def _window_span(extent: int, spacing: int) -> int: return spacing * (extent - 1) + 1 -def _window_indices( - X: TensorVariable, kernel_size: Sequence[int], stride: Sequence[int], dilation: Sequence[int] -) -> list[TensorVariable]: - """One advanced index per spatial axis, carrying that axis's windows and its taps. - - Each broadcasts against the others, so indexing with all of them at once gives windows-then-taps in - order. Shared by the gather's reference graph and by the scatter that reverses it. - """ - n_spatial = len(kernel_size) - indices = [] - for axis, (extent, step, spacing) in enumerate(zip(kernel_size, stride, dilation)): - span = _window_span(extent, spacing) - starts = pt.arange(0, X.shape[1 + axis] - span + 1, step) - window = starts[:, None] + pt.arange(extent)[None, :] * spacing - - pattern: list[int | str] = ["x"] * (2 * n_spatial) - pattern[axis] = 0 - pattern[n_spatial + axis] = 1 - indices.append(window.dimshuffle(*pattern)) - return indices - - -def _scatter_patches( - cotangent: TensorVariable, - X: TensorVariable, - kernel_size: Sequence[int], - stride: Sequence[int], - dilation: Sequence[int], -) -> TensorVariable: - """Add each window's cotangent back at the position it was gathered from. - - Windows overlap, so a position reached by several of them accumulates all of their contributions -- - which is why this is a scatter-add rather than an assignment. - """ - indices = _window_indices(X, kernel_size, stride, dilation) - zeros = pt.zeros(X.shape, dtype=cotangent.dtype) - return pt.inc_subtensor(zeros[(slice(None), *indices, slice(None))], cotangent) - - class Im2Col(Op): """ Gather every window a kernel visits, as one node a backend can put a real copy behind. @@ -284,36 +245,6 @@ def pullback(self, inputs, outputs, cotangents): return [gathered, *(disconnected_type() for _ in self.kernel_size)] -def _extract_patches( - X: TensorVariable, - kernel_size: Sequence[int], - stride: Sequence[int], - dilation: Sequence[int], -) -> TensorVariable: - """ - Gather every window a kernel visits, for an input with any number of spatial axes. - - Parameters - ---------- - X : TensorVariable - Input of shape ``(batch, *spatial, channels)``. - kernel_size : sequence of int - Window extent along each spatial axis. Its length is the number of spatial axes. - stride : sequence of int - Step between windows along each spatial axis. - dilation : sequence of int - Spacing between the taps of one window along each spatial axis. - - Returns - ------- - patches : TensorVariable - Shape ``(batch, *out_spatial, *kernel_size, channels)``, where ``out_spatial`` counts the - windows that fit. - """ - indices = _window_indices(X, kernel_size, stride, dilation) - return X[(slice(None), *indices, slice(None))] - - class ConvLayer(UnaryLayerOp): """ Cross-correlate an input with a kernel, over any number of spatial axes. From bcab78fe7df8f5b12d020f7f2bdb47be2f4c69e2 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:16:05 -0500 Subject: [PATCH 07/10] Trim the transposed-convolution tests --- tests/test_conv_transpose.py | 124 ++++++++++++++++------------------- 1 file changed, 56 insertions(+), 68 deletions(-) diff --git a/tests/test_conv_transpose.py b/tests/test_conv_transpose.py index 5e012ed..8d2f456 100644 --- a/tests/test_conv_transpose.py +++ b/tests/test_conv_transpose.py @@ -3,17 +3,16 @@ import pytensor.tensor as pt import pytest -from pytensor_ml.layers import Conv2D, ConvTranspose1D, ConvTranspose2D, Input, Linear -from pytensor_ml.layers.combinators import Flatten -from pytensor_ml.loss import SquaredError -from pytensor_ml.model import Model -from pytensor_ml.optim import adam +from pytensor.gradient import verify_grad +from pytensor.graph import graph_replace + +from pytensor_ml.layers import Conv2D, ConvTranspose1D, ConvTranspose2D floatX = pytensor.config.floatX ATOL = 1e-4 if floatX == "float32" else 1e-8 -@pytest.fixture(scope="module") +@pytest.fixture def rng(): return np.random.default_rng(sum(map(ord, "conv transpose"))) @@ -45,22 +44,21 @@ def scatter_through_kernel(X, W, stride, padding, output_padding, dilation): @pytest.mark.parametrize( - "stride, padding, output_padding, dilation", + "stride, padding, output_padding, dilation, bias", [ - (1, 0, 0, 1), - (2, 0, 0, 1), - (2, 1, 1, 1), - (3, 2, 2, 1), - (1, 1, 0, 2), - ((2, 3), (1, 2), (1, 2), 1), + (1, 0, 0, 1, False), + (2, 1, 1, 1, True), + ((2, 3), (1, 2), (1, 2), (2, 1), False), ], - ids=["plain", "strided", "strided_cropped", "wide_stride", "dilated", "per_axis"], + ids=["plain", "strided_cropped", "per_axis"], ) def test_conv_transpose_2d_scatters_its_input_through_the_kernel( - stride, padding, output_padding, dilation, rng + stride, padding, output_padding, dilation, bias, rng ): """Every argument changes where a contribution lands, so each is exercised against the scatter the - layer stands for. The per-axis case is the one that catches an axis order swapped somewhere.""" + layer stands for. The per-axis case is the one that catches an axis order swapped somewhere. The + bias is added after the scatter, so a layer that folded it in before would add it once per + contributing tap instead of once per output position.""" X_np = rng.normal(size=(2, 5, 6, 3)).astype(floatX) X = pt.tensor("X", shape=(None, 5, 6, 3), dtype=floatX) layer = ConvTranspose2D( @@ -72,9 +70,13 @@ def test_conv_transpose_2d_scatters_its_input_through_the_kernel( padding=padding, output_padding=output_padding, dilation=dilation, - bias=False, + bias=bias, ) layer.W.set_value(rng.normal(size=layer.W.get_value().shape).astype(floatX)) + expected_bias = 0.0 + if bias: + expected_bias = rng.normal(size=(4,)).astype(floatX) + layer.b.set_value(expected_bias) expected = scatter_through_kernel( X_np, @@ -84,17 +86,12 @@ def test_conv_transpose_2d_scatters_its_input_through_the_kernel( layer.output_padding, layer.dilation, ) - np.testing.assert_allclose(pytensor.function([X], layer(X))(X_np), expected, atol=ATOL) + np.testing.assert_allclose( + pytensor.function([X], layer(X))(X_np), expected + expected_bias, atol=ATOL + ) -@pytest.mark.parametrize( - "stride, padding, output_padding", - [(1, 0, 0), (2, 1, 1), (3, 1, 2)], - ids=["plain", "strided_cropped", "wide_stride"], -) -def test_conv_transpose_1d_scatters_its_input_through_the_kernel( - stride, padding, output_padding, rng -): +def test_conv_transpose_1d_scatters_its_input_through_the_kernel(rng): """The rank-1 layer shares `_ConvTransposeNd`, so this checks the spatial count reaches every argument rather than repeating the rank-2 coverage.""" X_np = rng.normal(size=(2, 7, 3)).astype(floatX) @@ -104,9 +101,9 @@ def test_conv_transpose_1d_scatters_its_input_through_the_kernel( in_channels=3, out_channels=5, kernel_size=4, - stride=stride, - padding=padding, - output_padding=output_padding, + stride=2, + padding=1, + output_padding=1, bias=False, ) layer.W.set_value(rng.normal(size=layer.W.get_value().shape).astype(floatX)) @@ -122,24 +119,6 @@ def test_conv_transpose_1d_scatters_its_input_through_the_kernel( np.testing.assert_allclose(pytensor.function([X], layer(X))(X_np), expected, atol=ATOL) -def test_the_bias_is_added_once_per_output_channel(rng): - """The bias is added after the scatter, so a layer that folded it in before would add it once per - contributing tap instead of once per output position.""" - X_np = rng.normal(size=(2, 5, 5, 3)).astype(floatX) - X = pt.tensor("X", shape=(None, 5, 5, 3), dtype=floatX) - layer = ConvTranspose2D("transpose", in_channels=3, out_channels=4, kernel_size=3, stride=2) - layer.W.set_value(rng.normal(size=layer.W.get_value().shape).astype(floatX)) - bias = rng.normal(size=(4,)).astype(floatX) - layer.b.set_value(bias) - - with_bias = pytensor.function([X], layer(X))(X_np) - layer.b.set_value(np.zeros_like(bias)) - without = pytensor.function([X], layer(X))(X_np) - np.testing.assert_allclose( - with_bias - without, np.broadcast_to(bias, with_bias.shape), atol=ATOL - ) - - def test_conv_transpose_inverts_a_convolutions_shape(): """The operation exists to undo what a convolution did to a shape, so a transpose built with the same geometry has to return the size the convolution consumed.""" @@ -175,23 +154,26 @@ def test_padding_does_not_take_the_forward_layers_keywords(): ConvTranspose2D("t", in_channels=1, out_channels=1, kernel_size=3, padding="same") -def test_an_encoder_decoder_trains_end_to_end(rng): - """The step the layer exists for: something that downsamples, then upsamples back, learning through - both. Gradients have to cross the transpose, which is the op run forward, so this is what says its - pullback works inside a real graph rather than only under `verify_grad`.""" - X = Input("X", shape=(None, 8, 8, 2)) - encoded = Conv2D("down", in_channels=2, out_channels=4, kernel_size=3, stride=2)(X) - decoded = ConvTranspose2D("up", in_channels=4, out_channels=2, kernel_size=3, stride=2)(encoded) - y = Linear("head", n_in=7 * 7 * 2, n_out=1)(Flatten(decoded)) - - model = Model(y).initialize(seed=1) - step = model.compile_train(adam(learning_rate=0.05), SquaredError()) +def test_the_gradient_crosses_the_output_padding_and_the_crop(rng): + """The pullback has to undo the leftover row `output_padding` adds and the rows the crop removes, + and no forward comparison sees either of those through the gradient.""" + layer = ConvTranspose2D( + "transpose", + in_channels=2, + out_channels=3, + kernel_size=3, + stride=2, + padding=1, + output_padding=1, + bias=False, + ) + X_np = rng.normal(size=(2, 4, 3, 2)).astype(floatX) + W_np = rng.normal(size=layer.W.get_value().shape).astype(floatX) - X_np = rng.normal(size=(32, 8, 8, 2)).astype(floatX) - y_np = X_np.sum(axis=(1, 2, 3))[:, None].astype(floatX) + def transpose(X, W): + return graph_replace(layer(X), {layer.W: W}) - losses = [float(step(X_np, y_np)) for _ in range(50)] - assert losses[-1] < losses[0] / 5 + verify_grad(transpose, [X_np, W_np], rng=np.random.default_rng(0)) def test_cropping_more_than_the_output_holds_is_rejected(): @@ -216,16 +198,22 @@ def test_a_crop_that_only_a_short_input_cannot_afford_is_allowed(): def test_a_symbolic_spatial_extent_skips_the_crop_check(rng): """Only a statically known length can be checked, so a symbolic one is let through rather than - guessed at -- and then has to compute the same answer a declared length would.""" - symbolic = pt.tensor("symbolic", shape=(None, None, None, 3), dtype=floatX) - declared = pt.tensor("declared", shape=(2, 5, 5, 3), dtype=floatX) + guessed at -- and then has to compute the scatter a declared length would.""" + X = pt.tensor("X", shape=(None, None, None, 3), dtype=floatX) layer = ConvTranspose2D( "t", in_channels=3, out_channels=4, kernel_size=3, padding=2, bias=False ) layer.W.set_value(rng.normal(size=layer.W.get_value().shape).astype(floatX)) X_np = rng.normal(size=(2, 5, 5, 3)).astype(floatX) - unchecked = pytensor.function([symbolic], layer(symbolic))(X_np) - checked = pytensor.function([declared], layer(declared))(X_np) + unchecked = pytensor.function([X], layer(X))(X_np) + expected = scatter_through_kernel( + X_np, + layer.W.get_value(), + layer.stride, + [before for before, _ in layer.padding], + layer.output_padding, + layer.dilation, + ) assert unchecked.shape == (2, 3, 3, 4) - np.testing.assert_allclose(unchecked, checked, atol=ATOL) + np.testing.assert_allclose(unchecked, expected, atol=ATOL) From 1d56dde6b88ba30d92ca21364a37cb70a3a60518 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:16:07 -0500 Subject: [PATCH 08/10] Trim the padding tests --- tests/test_padding.py | 48 +++++++++++-------------------------------- 1 file changed, 12 insertions(+), 36 deletions(-) diff --git a/tests/test_padding.py b/tests/test_padding.py index 3c545cc..e799063 100644 --- a/tests/test_padding.py +++ b/tests/test_padding.py @@ -5,19 +5,17 @@ from pytensor_ml.layers import ( ConstantPad1D, - ConstantPad2D, ReflectionPad1D, ReflectionPad2D, ReplicationPad1D, ReplicationPad2D, - ZeroPad1D, ZeroPad2D, ) floatX = pytensor.config.floatX -@pytest.fixture(scope="module") +@pytest.fixture def rng(): return np.random.default_rng(sum(map(ord, "padding"))) @@ -31,16 +29,10 @@ def rng(): "constant", {"constant_values": 0.0}, ), - ( - ConstantPad2D(padding=((2, 1), (1, 3)), value=-1.5), - ((2, 1), (1, 3)), - "constant", - {"constant_values": -1.5}, - ), (ReflectionPad2D(padding=((1, 2), (2, 1))), ((1, 2), (2, 1)), "reflect", {}), (ReplicationPad2D(padding=((2, 1), (1, 2))), ((2, 1), (1, 2)), "edge", {}), ], - ids=["zero", "constant", "reflection", "replication"], + ids=["zero", "reflection", "replication"], ) def test_each_mode_pads_the_spatial_axes_like_numpy( layer, pad_width, numpy_mode, numpy_kwargs, rng @@ -59,12 +51,10 @@ def test_each_mode_pads_the_spatial_axes_like_numpy( @pytest.mark.parametrize( "layer, numpy_mode, numpy_kwargs", [ - (ZeroPad1D(padding=(2, 4)), "constant", {"constant_values": 0.0}), (ConstantPad1D(padding=(2, 4), value=3.5), "constant", {"constant_values": 3.5}), - (ReflectionPad1D(padding=(2, 4)), "reflect", {}), (ReplicationPad1D(padding=(2, 4)), "edge", {}), ], - ids=["zero", "constant", "reflection", "replication"], + ids=["constant", "replication"], ) def test_one_spatial_axis_reads_a_bare_pair_as_its_two_ends(layer, numpy_mode, numpy_kwargs, rng): """Over a single axis `(2, 4)` can only mean its two ends, and that is how torch reads it too -- @@ -87,28 +77,6 @@ def test_padding_wider_than_the_axis_keeps_reflecting(rng): np.testing.assert_allclose(padded, np.pad(X_np, [(0, 0), (5, 5), (0, 0)], mode="reflect")) -def test_the_batch_and_channel_axes_are_left_alone(rng): - """Padding is spatial. A layer that padded every axis would still produce a plausible-looking - array, so the untouched axes are asserted rather than assumed.""" - X_np = rng.normal(size=(2, 5, 6, 3)).astype(floatX) - X = pt.tensor("X", shape=(None, 5, 6, 3), dtype=floatX) - - padded = pytensor.function([X], ZeroPad2D(padding=2)(X))(X_np) - assert padded.shape == (2, 9, 10, 3) - np.testing.assert_allclose(padded[:, 2:-2, 2:-2, :], X_np) - - -def test_gradients_reach_the_input_through_the_padding(rng): - """The padded positions are constants, so the gradient of a padded output has to arrive back at - exactly the elements that came from the input and nowhere else.""" - X_np = rng.normal(size=(1, 4, 4, 1)).astype(floatX) - X = pt.tensor("X", shape=(1, 4, 4, 1), dtype=floatX) - - padded = ZeroPad2D(padding=1)(X) - gradient = pytensor.function([X], pt.grad(padded.sum(), X))(X_np) - np.testing.assert_allclose(gradient, np.ones_like(X_np)) - - def test_an_input_of_the_wrong_rank_is_rejected(): """The layers are channels-last over a fixed number of spatial axes, so a channels-first image or a batch of vectors is a mistake worth naming rather than padding the wrong axes.""" @@ -122,11 +90,19 @@ def test_an_input_of_the_wrong_rank_is_rejected(): (-1, "cannot be negative"), (((1, 2), (3, -4)), "cannot be negative"), ((1, 2, 3), "one amount per axis"), + (((1, 2, 3), (1, 1)), r"is a \(before, after\) pair"), ], - ids=["negative_scalar", "negative_in_pair", "wrong_count"], + ids=["negative_scalar", "negative_in_pair", "wrong_count", "triple_for_one_axis"], ) def test_padding_amounts_are_validated(padding, message): """Each of these describes something the layer cannot do, and each would otherwise surface as a confusing shape further downstream.""" with pytest.raises(ValueError, match=message): ZeroPad2D(padding=padding) + + +def test_one_amount_per_axis_pads_both_of_its_sides(): + """Over two axes a flat pair is one amount per axis, so `(1, 2)` puts 1 on the top and bottom and + 2 on the left and right. The 1-D layers read the same pair as one axis's two ends, and the two + readings must not leak into each other.""" + assert ZeroPad2D(padding=(1, 2)).padding == ((1, 1), (2, 2)) From 4ae6211c1c356a5c4a52ca67417399d3900c5167 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:16:09 -0500 Subject: [PATCH 09/10] Trim the pooling tests --- tests/test_pool.py | 100 +++++++++++++++++++-------------------------- 1 file changed, 43 insertions(+), 57 deletions(-) diff --git a/tests/test_pool.py b/tests/test_pool.py index 09022e2..c490194 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -63,24 +63,23 @@ def test_average_pooling_counts_padding_as_zero(): np.testing.assert_allclose(pooled.eval({X: X_np})[0, :, 0], [-1.5, -1.5], atol=ATOL) -def test_pooling_steps_by_the_kernel_unless_told_otherwise(rng): +def test_pooling_steps_by_the_kernel_unless_told_otherwise(): """Pooling tiles by default where convolution slides, so the default stride is the kernel extent. Anyone reading the two signatures together will take that for a mistake, so it is pinned here.""" X = pt.tensor("X", shape=(2, 8, 8, 3)) - X_np = rng.normal(size=(2, 8, 8, 3)).astype(floatX) - tiled = MaxPool2D("tiled", kernel_size=2)(X) - slid = MaxPool2D("slid", kernel_size=2, stride=1)(X) - - assert tiled.eval({X: X_np}).shape == (2, 4, 4, 3) - assert slid.eval({X: X_np}).shape == (2, 7, 7, 3) + assert MaxPool2D("tiled", kernel_size=2)(X).type.shape == (2, 4, 4, 3) + assert MaxPool2D("slid", kernel_size=2, stride=1)(X).type.shape == (2, 7, 7, 3) def test_max_pooling_routes_its_gradient_to_the_window_maximum(): """A max pool's pullback sends the whole cotangent to the position that won and nothing anywhere - else, so the gradient of a summed pool is an indicator of the per-window argmax.""" + else, so the gradient of a summed pool is an indicator of the per-window argmax. Max pooling + usually follows a rectifier, so a window clamped entirely to zero is routine rather than a + measure-zero curiosity. Crediting every tied tap would return more gradient than the window + received, so the earliest one takes it, as jax, torch and mlx all do.""" X = pt.tensor("X", shape=(1, 6, 1)) - X_np = np.array([3.0, 1.0, 2.0, 5.0, 4.0, 0.0], dtype=floatX).reshape(1, 6, 1) + X_np = np.array([3.0, 1.0, 2.0, 5.0, 1.0, 1.0], dtype=floatX).reshape(1, 6, 1) gradient = pt.grad(MaxPool1D("pool", kernel_size=2)(X).sum(), X) @@ -91,13 +90,31 @@ def test_max_pooling_routes_its_gradient_to_the_window_maximum(): def test_average_pooling_spreads_its_gradient_over_the_window(rng): """The counterpart to the max routing: every position in a window contributed equally, so each - receives the same share rather than one taking all of it.""" + receives the same share rather than one taking all of it. The windows overlap, so a position + collects one share from every window that covers it, and a pullback that wrote its shares + instead of adding them would keep only the last.""" X = pt.tensor("X", shape=(1, 6, 1)) X_np = rng.normal(size=(1, 6, 1)).astype(floatX) - gradient = pt.grad(AvgPool1D("pool", kernel_size=3)(X).sum(), X) + gradient = pt.grad(AvgPool1D("pool", kernel_size=3, stride=1)(X).sum(), X) + + np.testing.assert_allclose( + gradient.eval({X: X_np})[0, :, 0], np.array([1, 2, 3, 3, 2, 1]) / 3, atol=ATOL + ) + + +def test_dilated_pooling_reads_every_other_element(): + """A dilation of 2 spaces a window's taps two apart, so with a stride of 1 the windows are + `(0, 2), (1, 3), (2, 4), (3, 5)`. The stride and dilation differ, so a kernel that swapped them + reads other windows, and one that ignored the dilation reads adjacent pairs.""" + X = pt.tensor("X", shape=(1, 6, 1)) + X_np = np.array([3.0, 0.0, 1.0, 4.0, 2.0, 5.0], dtype=floatX).reshape(1, 6, 1) - np.testing.assert_allclose(gradient.eval({X: X_np})[0, :, 0], np.full(6, 1 / 3), atol=ATOL) + pooled = MaxPool1D("pool", kernel_size=2, stride=1, dilation=2)(X) + values, gradient = pytensor.function([X], [pooled, pt.grad(pooled.sum(), X)])(X_np) + + np.testing.assert_allclose(values[0, :, 0], [3.0, 4.0, 2.0, 5.0], atol=ATOL) + np.testing.assert_allclose(gradient[0, :, 0], [1.0, 0.0, 0.0, 1.0, 1.0, 1.0], atol=ATOL) @pytest.mark.parametrize( @@ -118,13 +135,6 @@ def test_pooling_over_one_axis_matches_a_degenerate_second(layers, rng): ) -def test_pooling_keeps_the_channel_axis_and_the_shape_it_can_work_out(): - """Pooling never mixes channels, and a layer downstream has to size itself from the graph.""" - X = pt.tensor("X", shape=(32, 28, 28, 16)) - - assert MaxPool2D("pool", kernel_size=2)(X).type.shape == (32, 14, 14, 16) - - def test_pooling_rejects_an_input_too_small_for_one_window(): """Shared with the convolution layers: a window that does not fit yields no windows at all, and an empty answer is worse than an error.""" @@ -134,53 +144,29 @@ def test_pooling_rejects_an_input_too_small_for_one_window(): MaxPool2D("pool", kernel_size=4)(X) -def test_max_pooling_gives_a_tied_window_to_one_tap(): - """Max pooling usually follows a rectifier, so a window clamped entirely to zero is routine rather - than a measure-zero curiosity. Crediting every tied tap would return more gradient than the window - received, so the earliest one takes it, as jax, torch and mlx all do.""" - X = pt.tensor("X", shape=(1, 4, 1)) - X_np = np.array([1.0, 1.0, 2.0, 2.0], dtype=floatX).reshape(1, 4, 1) - - gradient = pt.grad(MaxPool1D("pool", kernel_size=2)(X).sum(), X) - - np.testing.assert_allclose(gradient.eval({X: X_np})[0, :, 0], [1.0, 0.0, 1.0, 0.0], atol=ATOL) - - # No linker runs the pooling ops' inner graph once every backend dispatches them, so it is reached # deliberately here. It is the portable definition of what pooling means, and a backend's kernel is # only correct if it agrees with it. FALLBACK = Mode(linker="py", optimizer="fast_compile") -@pytest.mark.parametrize("padding", ["valid", "same"], ids=["valid", "same"]) -@pytest.mark.parametrize("layer_cls", [MaxPool1D, AvgPool1D], ids=["max", "avg"]) -def test_the_portable_graph_agrees_with_the_kernel_that_replaces_it(layer_cls, padding, rng): +@pytest.mark.parametrize("layer_cls", [MaxPool2D, AvgPool2D], ids=["max", "avg"]) +def test_the_portable_graph_agrees_with_the_kernel_that_replaces_it(layer_cls, rng): """Every backend dispatches past the inner graph, so nothing else exercises it -- and it is what runs wherever a kernel is missing. A kernel and a fallback that disagree is a result that changes - with the backend.""" - X_np = rng.normal(size=(2, 9, 3)).astype(floatX) + with the backend. Every geometry argument differs per axis, and the input is drawn from three + values so that windows tie: `pt.max` would hand the cotangent to every tied tap, and the fallback + is where that would go unnoticed.""" + X_np = rng.integers(0, 3, size=(2, 7, 6, 3)).astype(floatX) X = pt.tensor("X", shape=X_np.shape) - out = layer_cls("pool", kernel_size=3, stride=2, padding=padding)(X) - - np.testing.assert_allclose( - pytensor.function([X], out, mode=FALLBACK)(X_np), - pytensor.function([X], out)(X_np), - atol=ATOL, - ) - - -def test_the_portable_graph_routes_a_tied_window_the_way_the_kernel_does(): - """The tie rule is the part of the contract most easily lost: `pt.max` would hand the cotangent to - every tied tap, and the fallback is where that would go unnoticed.""" - X_np = np.array([1.0, 1.0, 2.0, 2.0], dtype=floatX).reshape(1, 4, 1) - X = pt.tensor("X", shape=X_np.shape) - gradient = pt.grad(MaxPool1D("pool", kernel_size=2)(X).sum(), X) - - np.testing.assert_allclose( - pytensor.function([X], gradient, mode=FALLBACK)(X_np)[0, :, 0], - [1.0, 0.0, 1.0, 0.0], - atol=ATOL, - ) + out = layer_cls("pool", kernel_size=(3, 2), stride=(2, 1), dilation=(1, 2), padding="same")(X) + outputs = [out, pt.grad(out.sum(), X)] + + # The py linker still builds C thunks unless told not to, and those compiles dominate the cost. + with pytensor.config.change_flags(cxx=""): + fallback = pytensor.function([X], outputs, mode=FALLBACK)(X_np) + for portable, kernel in zip(fallback, pytensor.function([X], outputs)(X_np)): + np.testing.assert_allclose(portable, kernel, atol=ATOL) def test_the_pooling_gradient_op_rejects_a_reduction_it_does_not_have(): From 1103b945c0b1e066a60e0dd39829b9445fb9f413 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Fri, 2 Oct 2026 15:16:11 -0500 Subject: [PATCH 10/10] Trim the upsampling tests --- tests/test_upsample.py | 140 +++++++++++++++++------------------------ 1 file changed, 57 insertions(+), 83 deletions(-) diff --git a/tests/test_upsample.py b/tests/test_upsample.py index fdcadd1..bec1418 100644 --- a/tests/test_upsample.py +++ b/tests/test_upsample.py @@ -3,8 +3,6 @@ import pytensor.tensor as pt import pytest -from pytensor.gradient import verify_grad - from pytensor_ml.layers import Upsample1D, Upsample2D floatX = pytensor.config.floatX @@ -16,14 +14,12 @@ def rng(): return np.random.default_rng(sum(map(ord, "pytensor_ml upsample"))) -@pytest.mark.parametrize("factor", [2, 3, (2, 3)], ids=["square", "square_odd", "non_square"]) -def test_nearest_repeats_each_element(factor, rng): - height_factor, width_factor = (factor, factor) if isinstance(factor, int) else factor +def test_nearest_repeats_each_element(rng): X = pt.tensor("X", shape=(None, 3, 4, 2)) - out = Upsample2D(scale_factor=factor)(X) + out = Upsample2D(scale_factor=(2, 3))(X) X_np = rng.normal(size=(2, 3, 4, 2)).astype(floatX) - expected = np.repeat(np.repeat(X_np, height_factor, axis=1), width_factor, axis=2) + expected = np.repeat(np.repeat(X_np, 2, axis=1), 3, axis=2) np.testing.assert_allclose(out.eval({X: X_np}), expected, rtol=1e-6, atol=ATOL) @@ -39,6 +35,15 @@ def test_nearest_indexes_by_floor_of_the_ratio(): np.testing.assert_allclose(out.eval({X: X_np}).ravel(), [0, 0, 1, 2, 2, 3, 4]) +HALF_PIXEL_INPUT = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=floatX).reshape(1, 2, 2, 1) +HALF_PIXEL_DOUBLED = [ + [1.00, 1.25, 1.75, 2.00], + [1.50, 1.75, 2.25, 2.50], + [2.50, 2.75, 3.25, 3.50], + [3.00, 3.25, 3.75, 4.00], +] + + def test_bilinear_half_pixel_convention(): """With align_corners=False each element covers a unit interval and output centers map to input centers, so a 2x2 doubled leaves the corners unchanged and the outermost samples land outside @@ -46,15 +51,9 @@ def test_bilinear_half_pixel_convention(): X = pt.tensor("X", shape=(1, 2, 2, 1)) out = Upsample2D(scale_factor=2, mode="bilinear")(X) - X_np = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=floatX).reshape(1, 2, 2, 1) - expected = [ - [1.00, 1.25, 1.75, 2.00], - [1.50, 1.75, 2.25, 2.50], - [2.50, 2.75, 3.25, 3.50], - [3.00, 3.25, 3.75, 4.00], - ] - - np.testing.assert_allclose(out.eval({X: X_np})[0, :, :, 0], expected, rtol=1e-6, atol=ATOL) + np.testing.assert_allclose( + out.eval({X: HALF_PIXEL_INPUT})[0, :, :, 0], HALF_PIXEL_DOUBLED, rtol=1e-6, atol=ATOL + ) def test_bilinear_corner_aligned_convention(): @@ -76,60 +75,52 @@ def test_bilinear_corner_aligned_convention(): np.testing.assert_allclose(out.eval({X: X_np})[0, :, :, 0], expected, rtol=1e-6, atol=ATOL) -@pytest.mark.parametrize( - "align_corners, expected", - [(False, [1.0, 1.25, 1.75, 2.25, 2.75, 3.0]), (True, [1.0, 1.4, 1.8, 2.2, 2.6, 3.0])], - ids=["half_pixel", "corner_aligned"], -) -def test_linear_interpolates_along_a_sequence(align_corners, expected): +def test_linear_interpolates_along_a_sequence(): X = pt.tensor("X", shape=(1, 3, 1)) - out = Upsample1D(scale_factor=2, mode="linear", align_corners=align_corners)(X) + out = Upsample1D(scale_factor=2, mode="linear")(X) X_np = np.array([1.0, 2.0, 3.0], dtype=floatX).reshape(1, 3, 1) + expected = [1.0, 1.25, 1.75, 2.25, 2.75, 3.0] np.testing.assert_allclose(out.eval({X: X_np}).ravel(), expected, rtol=1e-6, atol=ATOL) -@pytest.mark.parametrize("size", [(8, 6), (2, 9)], ids=["grow", "shrink_one_axis"]) -def test_explicit_size_resamples_to_that_extent(size, rng): - """A size that is not a uniform factor of the input, so nothing else covers these extents. The - axes are given height-first, which a swap turns into the wrong shape rather than wrong values.""" - X = pt.tensor("X", shape=(None, 3, 4, 2)) - out = Upsample2D(size=size, mode="bilinear")(X) - - res = out.eval({X: rng.normal(size=(2, 3, 4, 2)).astype(floatX)}) - - assert out.type.shape == (None, *size, 2) - assert res.shape == (2, *size, 2) +def half_pixel_interpolation(X, size): + """Separable linear interpolation where output centers map to input centers. `np.interp` holds the + end values beyond the outermost centers, which is the clamp the half-pixel convention needs.""" + out = X + for axis, out_extent in enumerate(size, start=1): + in_extent = out.shape[axis] + source = (np.arange(out_extent) + 0.5) * in_extent / out_extent - 0.5 + out = np.apply_along_axis( + lambda values: np.interp(source, np.arange(in_extent), values), axis, out + ) + return out -@pytest.mark.parametrize("mode", ["nearest", "bilinear"]) -def test_explicit_size_agrees_with_the_equivalent_factor(mode, rng): - """Asking a 3x4 input for size (6, 8) is the same resampling as scale_factor=2, reached through - the other branch of the extent resolution.""" +def test_explicit_size_resamples_to_that_extent(rng): + """A size that is not a uniform factor of the input, shrinking one axis and growing the other, so + nothing else covers these extents. The axes are given height-first, which a swap turns into the + wrong shape rather than wrong values.""" X = pt.tensor("X", shape=(None, 3, 4, 2)) + out = Upsample2D(size=(2, 9), mode="bilinear")(X) X_np = rng.normal(size=(2, 3, 4, 2)).astype(floatX) + assert out.type.shape == (None, 2, 9, 2) np.testing.assert_allclose( - Upsample2D(size=(6, 8), mode=mode)(X).eval({X: X_np}), - Upsample2D(scale_factor=2, mode=mode)(X).eval({X: X_np}), - rtol=1e-6, - atol=ATOL, + out.eval({X: X_np}), half_pixel_interpolation(X_np, (2, 9)), rtol=1e-6, atol=ATOL ) -@pytest.mark.parametrize( - "align_corners, expected", [(False, 6.5), (True, 1.0)], ids=["half_pixel", "corner_aligned"] -) -def test_a_single_element_output_is_well_defined(align_corners, expected): +def test_a_single_element_output_is_well_defined(): """Spreading over a closed interval divides by one less than the output extent, which is zero when the output is a single element.""" X = pt.tensor("X", shape=(1, 3, 4, 1)) - out = Upsample2D(size=1, mode="bilinear", align_corners=align_corners)(X) + out = Upsample2D(size=1, mode="bilinear", align_corners=True)(X) X_np = np.arange(1, 13, dtype=floatX).reshape(1, 3, 4, 1) - np.testing.assert_allclose(out.eval({X: X_np}).ravel(), [expected], rtol=1e-6, atol=ATOL) + np.testing.assert_allclose(out.eval({X: X_np}).ravel(), [1.0], rtol=1e-6, atol=ATOL) def test_a_known_extent_survives_the_gather(): @@ -140,47 +131,30 @@ def test_a_known_extent_survives_the_gather(): assert Upsample2D(scale_factor=2)(X).type.shape == (None, 6, 8, 2) -def test_an_unknown_extent_still_resamples(rng): +def test_an_unknown_extent_still_resamples(): """The extents go symbolic when the input's are unknown, which must change the graph and not the - answer.""" - dynamic = pt.tensor("dynamic", shape=(None, None, None, 2)) - static = pt.tensor("static", shape=(None, 3, 4, 2)) - out = Upsample2D(scale_factor=2, mode="bilinear")(dynamic) - - assert out.type.shape == (None, None, None, 2) + answer. A square input cannot tell the two spatial axes apart, so a non-square one is fed to the + same function to pin which extent each axis doubles.""" + X = pt.tensor("X", shape=(None, None, None, None)) + out = Upsample2D(scale_factor=2, mode="bilinear")(X) + resample = pytensor.function([X], out) - X_np = rng.normal(size=(1, 3, 4, 2)).astype(floatX) + assert out.type.shape == (None, None, None, None) np.testing.assert_allclose( - out.eval({dynamic: X_np}), - Upsample2D(scale_factor=2, mode="bilinear")(static).eval({static: X_np}), - rtol=1e-6, - atol=ATOL, + resample(HALF_PIXEL_INPUT)[0, :, :, 0], HALF_PIXEL_DOUBLED, rtol=1e-6, atol=ATOL ) + assert resample(np.zeros((1, 3, 4, 2), dtype=floatX)).shape == (1, 6, 8, 2) -@pytest.mark.parametrize("mode", ["nearest", "bilinear"]) -def test_upsample_gradient(mode, rng): - X_np = rng.normal(size=(2, 3, 4, 2)).astype(floatX) - - verify_grad(Upsample2D(scale_factor=2, mode=mode), [X_np], rng=np.random.default_rng(0)) - - -@pytest.mark.parametrize("mode", ["nearest", "bilinear"]) -@pytest.mark.parametrize("dtype", ["float32", "float64"]) -def test_resampling_keeps_the_input_dtype(dtype, mode): - """The source coordinates divide integer extents, and an integer division is float64 whatever - floatX is, so the upcast lands in a float32 network's activations.""" - with pytensor.config.change_flags(floatX=dtype): - X = pt.tensor("X", shape=(None, 3, 4, 2), dtype=dtype) - - assert Upsample2D(scale_factor=2, mode=mode)(X).dtype == dtype - - -@pytest.mark.parametrize("dtype", ["float16", "float32"]) -def test_bilinear_keeps_an_input_narrower_than_floatx(dtype): - """The interpolation weights are built from integer extents, and casting them to floatX would carry - a float32 or float16 input up to float64 under the default config.""" - with pytensor.config.change_flags(floatX="float64"): +@pytest.mark.parametrize( + "floatx, dtype", + [("float32", "float32"), ("float64", "float32"), ("float64", "float16")], + ids=["at_floatx", "float32_under_float64", "float16_under_float64"], +) +def test_bilinear_keeps_the_input_dtype(floatx, dtype): + """The interpolation weights are built from integer extents, and building them at floatX, or + dividing the extents as integers, would carry a narrower input up to float64.""" + with pytensor.config.change_flags(floatX=floatx): X = pt.tensor("X", shape=(None, 3, 4, 2), dtype=dtype) assert Upsample2D(scale_factor=2, mode="bilinear")(X).dtype == dtype