Skip to content

Commit 59b348e

Browse files
committed
Keep last POLAR_EXPRESS_COEFFICIENT unscaled
Apply safety_factor scaling to all POLAR_EXPRESS_COEFFICIENTS except the final tuple. The list comprehension now iterates over _unmodified_polar_express_coefficients[:-1] and the original last element is appended unchanged, preserving that coefficient (likely for correctness or numerical stability).
1 parent 237a0e8 commit 59b348e

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

modules/optimizer/muon.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@
4848
safety_factor = 1.05
4949
POLAR_EXPRESS_COEFFICIENTS = [
5050
(a / safety_factor , b / safety_factor**3 , c / safety_factor**5)
51-
for (a, b, c) in _unmodified_polar_express_coefficients
52-
]
51+
for (a, b, c) in _unmodified_polar_express_coefficients[: -1]
52+
] + [_unmodified_polar_express_coefficients[-1]]
5353

5454

5555
def zeropower_via_newtonschulz5(G: Tensor, steps: int, ns_coefficients: List[tuple]) -> Tensor:

0 commit comments

Comments
 (0)