A clock nothing advances compiles silently, pinning the schedule at step 0

[jessegrabowski]

function raises when the graph draws from a generator nothing advances, but the same mistake with a StepCounter compiles fine and every step reads schedule(0) -- with a warmup schedule that rate is exactly zero, so the network sits still and reports a flat loss with nothing in the logs.

import numpy as np
import pytensor.tensor as pt

from pytensor_ml.layers import Linear
from pytensor_ml.model import Model
from pytensor_ml.optim import adam, linear_schedule, scale
from pytensor_ml.params import step_counter
from pytensor_ml.pytensorf import function

x, y = pt.matrix("x"), pt.vector("y")
prediction = Linear("fc", 3, 1)(x)[:, 0]
model = Model(x, prediction).initialize(seed=0)
loss = pt.mean((prediction - y) ** 2)

clock = step_counter("schedule/step_count")
warmup = linear_schedule(0.0, total_steps=10, final_learning_rate=1.0)
updates = scale(warmup(clock))(adam(1e-1)(loss, model.weights), model.weights)

step = function([x, y], loss, updates=updates)  # nothing advances `clock`; compiles anyway

rng = np.random.default_rng(0)
features = rng.normal(size=(64, 3))
targets = features @ np.arange(3.0)
print([float(step(features, targets)) for _ in range(5)])  # constant: rate pinned at warmup(0) = 0
# fix: updates[clock] = clock.advance()