Function loss_fn

examples/python/linear_regression.py:26–27

(w)

def loss_fn(w):
    return 0.5 * mx.mean(mx.square(X @ w - y))


grad_fn = mx.grad(loss_fn)