Repository navigation
require_unique_state_names misses a duplicate name that is only read #122
Description
scalar_state and counter allocate a fresh shared variable on every call outside a reuses_state frame — the name is for the serialization boundary, not for lookup. Building the plateau scale the way the docstring reads, once for the policy and once for the rate, therefore produces two distinct variables under one name. The policy cuts its copy and the rate multiplies the other one. require_unique_state_names walks only the keys of the updates dict, and the rate-side copy is read but never written, so the guard cannot see it either.
import numpy as np
import pytensor.tensor as pt
from pytensor_ml.optim import adam, compile_train, reduce_on_plateau, scalar_state
from pytensor_ml.params import trainable
target = pt.scalar("target")
w = trainable(np.zeros(3), name="w")
loss = ((w - target) ** 2).sum() # starts at its minimum, so it never improves
rule = reduce_on_plateau(
adam(learning_rate=scalar_state("plateau/scale", fill_value=1.0) * 1e-3),
scalar_state("plateau/scale", fill_value=1.0),
patience=1,
factor=0.1,
)
step = compile_train(loss, rule, inputs=[target])
for _ in range(8):
step(0.0)
print([v.get_value() for v in step.get_shared() if v.name == "plateau/scale"])
# [array(1.), array(1.e-07)]
# the policy drove its copy to 1e-7; the one adam reads is still 1.0, and the rate never changedreduce_on_plateau(adam(1e-3), some_scale) — where the scale was never wired into the rate at all — is likewise a silent no-op. Checking that the scale is an ancestor of the rule's update expressions would catch both.
Activity
[-]Two scalar_state calls with one name give two variables, so reduce_on_plateau cuts a rate nothing reads[/-][+]require_unique_state_names misses a duplicate name that is only read[/+]jessegrabowski commentedon Sep 1, 2026
The documented example calls
scalar_stateonce and passes the result to both places. That avoids the collision entirely, so the bug as written isn't real. What's worth doing is wideningrequire_unique_state_namesto catch a duplicate name that's only read.