fix(muon): restore @compiler.compile() on muon_update
Now that the per-expert Python loop is replaced with a single batched
call to zeropower_via_newtonschulz5, muon_update has no dynamic control
flow and can be compiled again.
Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
