Tensorflow 2.10+ starts using optimizer.experimental_optimizer which requires a different update to the underlying optimizers. If/when that change is needed, you can find it in https://github.com/zanussbaum/mup-tf/commit/77ddc3a2cef7868709a3bbbfce21235ec512f2d9