scan_push_out_seq doesn't push a sequence-dependent Blockwise out of the loop

scan_push_out_seq only relocates Elemwise and DimShuffle, so a Blockwise applied per timestep — the input projection of any RNN, but equally any batched solve or cholesky — stays in the loop as N small ops instead of one batched one. (scan_push_out_dot1 is a different shape, a sit-sot accumulator read at [-1], and doesn't fire here.)

import timeit

import numpy as np
import pytensor
import pytensor.tensor as pt

from pytensor.scan import scan
from pytensor.scan.op import Scan

floatX = pytensor.config.floatX
W = pytensor.shared(np.random.default_rng(0).normal(size=(256, 256)).astype(floatX))
X = pt.tensor("X", shape=(200, 64, 256))
h0 = pt.zeros((64, 256), dtype=floatX)
X_np = np.random.default_rng(1).normal(size=(200, 64, 256)).astype(floatX)


def build(step, sequence=X):
    out = scan(step, sequences=[sequence], outputs_info=[h0], return_updates=False)
    return pytensor.function([X], out)


def inner_ops(fn):
    [node] = [n for n in fn.maker.fgraph.apply_nodes if isinstance(n.op, Scan)]
    return [str(n.op) for n in node.op.fgraph.apply_nodes]


def ms(fn):
    fn(X_np)
    return min(timeit.repeat(lambda: fn(X_np), number=5, repeat=5)) / 5 * 1e3


print(inner_ops(build(lambda x_t, h: pt.tanh(pt.exp(x_t) + h))))
# ['Composite{tanh((i0 + i1))}']         exp is gone, applied to the whole sequence at once

inside = build(lambda x_t, h: pt.tanh(x_t @ W + h))
print(inner_ops(inside))
# ['Dot', 'Composite{tanh((i0 + i1))}']  the dot stays, running 200 small matmuls

hoisted = build(lambda p_t, h: pt.tanh(p_t + h), sequence=X @ W)
print(f"{ms(inside):.1f} ms inside vs {ms(hoisted):.1f} ms hoisted by hand")
# ~24 ms inside vs ~21 ms hoisted by hand

Start by reproducing the example and reading scan_push_out_seq alongside scan_push_out_dot1, then inspect the Scan inner graph and its Blockwise nodes. Done means a sequence-dependent Blockwise operation is moved out of the loop as one batched operation, while the existing loop behavior remains correct and the reported inside-versus-hoisted performance gap is addressed.