Fix scale_by_rprop returning a one-step-stale update

[winklemad]

No linked issue — found by reading the code.

The bug

scale_by_rprop computes the correct current update in prev_updates (step_size * sign(g), zeroed where the gradient sign flips), but the emitted updates is built from state.prev_updates — the previous step's update — instead of the value just computed:

prev_updates = jax.tree.map(                       # correct current update
    lambda s, g, step_size: jnp.where(
        s < 0, jnp.zeros_like(g), step_size * jnp.sign(g)),
    sign, updates, step_sizes)
updates = jax.tree.map(
    lambda s, g, prev_g: jnp.where(s < 0, jnp.zeros_like(prev_g), prev_g),
    sign, prev_updates, state.prev_updates)         # <- returns prev_g (stale)

The freshly-computed prev_updates is passed in as g but the lambda never uses it — it returns prev_g (state.prev_updates). So every Rprop update lags one step behind, and the first step emits all zeros:

tx = optax.scale_by_rprop(learning_rate=0.1, eta_minus=0.5, eta_plus=1.2)
p = jnp.array([0.0]); s = tx.init(p)
for g in [1., 1., 1., -1., 1.]:
    u, s = tx.update(jnp.array([g]), s, p); print(float(u[0]))
# optax:   0.0, 0.1, 0.12, 0.0, 0.0
# correct: 0.1, 0.12, 0.144, 0.0, 0.072   (matches torch.optim.Rprop)

The generic convergence tests (alias_test.py) miss it because the optimizer still converges, just more slowly — they never assert per-step values.

Fix

Emit prev_updates (the current step's update) directly. state.prev_updates is still stored for the next step's sign computation — that use is correct and unchanged. With the fix the emitted sequence matches the reference Rprop exactly.

Testing

Added test_scale_by_rprop_step_values, asserting the per-step emitted updates against the hand-computed Rprop sequence. It fails before the change and passes after; optax/_src/transform_test.py passes (24 tests) and optax.rprop still converges. ruff is clean and the changed lines are pyink-formatted.