Split the layers CI group and cut its redundant tests

.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/

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
@@ -57,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.
@@ -283,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.
@@ -1305,16 +1237,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,14 +1270,20 @@ 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)
 
     # 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)

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.