Expose num_ns_steps

[janEbert]

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the
scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

[greptile-apps]

Greptile Summary

Exposes Muon’s construction-time num_ns_steps value through a read-only property.

• Stores the validated Newton–Schulz step count on each Muon instance.
• Prevents post-initialization assignment because the orthogonalization closure retains the original value.

Confidence Score: 5/5

The PR appears safe to merge with no actionable defects identified.

The new property is initialized from the same validated constructor argument used by the orthogonalization closure, and current subclass and serialization paths preserve that value consistently.

Important Files Changed

Filename	Overview
emerging_optimizers/orthogonalized_optimizers/muon.py	Adds a read-only property that consistently reports the same validated step count captured by Muon’s orthogonalization closure.

_{Reviews (1): Last reviewed commit: "Expose `num_ns_steps`" | Re-trigger Greptile}

[skyw]

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

It is very dangerous to modify scaled_orthogonalize_fn post initialization, so none of its arguments is maintained. What use cases do you have in mind, Jan? For simply use, inspect signiture can provide that information.

Having said that, we designed for relatively simple customization, like QKV split. If more complex customization is needed, we should revisit the entire design of OrthogonalizedOptimizer.