Function rms_norm_loop

benchmarks/python/rms_norm_bench.py:29–33

(g, x, w)

    mx.eval(x, w, y)

    def rms_norm_loop(g, x, w):
        gx, gw = x, w
        for _ in range(32):
            gx, gw = g(gx, gw, y)
        return gx, gw

    time_fn(rms_norm_loop, g1, x, w)
    time_fn(rms_norm_loop, g2, x, w)