StableMax/⊥Grad compatibility

[HydrogenBombaklot]

StableMax and ⊥Grad are new LLM training algorithms introduced in the Grokking at the Edge of Numerical Stability paper. Here is their repo.

These algorithms significantly boost generalization speed, resulting in faster, more efficient training and fine-tuning.

Can I drop-in these algorithms for MLX training/fine-tuning or do we need to integrate them more formally into MLX?

[angeloskath]

These look cool! I think the orthograd optimizer would be more suitable as a PR to https://github.com/stockeh/mlx-optimizers . Regarding stable max I think it should be fairly easy to add to any project as follows

@partial(mx.compile, shapeless=True)
def stablemax(x, axis=-1):
    x = mx.where(x < 0, -mx.log1p(-x), mx.log1p(x))
    return mx.softmax(x, axis=axis)

[HydrogenBombaklot]

Thank you kind stranger!