Stop Optim.step(scale_by) mutating the caller's gradients

[cemde]

Bug

Optim.step(scale_by=k) scaled each gradient with set(g, g * scale_by). set writes the new value into the Param object it is given, and that object is the caller's, not a copy. So the call mutated the gradient tree the caller still holds.

Impact

Stepping the same gradient tree twice compounds the scaling. With lr=1.0, scale_by=0.5, g=1.0, w0=1.0, two steps must give 1 - 2*(1.0*0.5*1.0) = 0.0; PCX gives 0.25, because the gradient is 0.5 on the second step and 0.25 on the third.

Every tutorial hits this: they compute one gradient tree and feed it to both optim_h and optim_w with scale_by=1.0/batch_size. The effective learning rate shrinks geometrically with nothing to show for it but slow convergence.

The same mechanism also compounds within a single step on a pxnn.shared model, since the outer tree_map visits an aliased Param once per reference.

Fix

return jtu.tree_map(lambda _g: _g * scale_by, g)

tree_map descends into the Param (it is a registered pytree), scales the value, and rebuilds a new Param via the registered unflatten. The caller's object is untouched.

Why tree_map and not type(g)(g * scale_by)

The registered unflatten copies __dict__ minus _value, so the subclass and all its attributes survive. Constructing through __init__ does not: a rebuilt VodeParam would lose frozen, the attribute the tutorials set on the output node and that M_hasnot(VodeParam, frozen=True) reads. A mask over the scaled tree would then select the wrong parameters.

OptimTree.step delegates here, so it inherits the fix.

Closes #67

[liukidar]

LGTM. Big fix :)