Gradients

value_and_grad

value_and_grad(f, m, args…) calls f(m, args…) and gives two results, its value and its gradient with respect to the params of m, so it’s bound to two names (E0411):

fn loss(m: Mlp, x: [b D], y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, x), y, over=C)

fn grads(m: Mlp, x: [b D], y: [b] Index[C]) -> Grad[Mlp]:
    l, g = value_and_grad(loss, m, x, y)
    return g

f is the name of a fn (not an extern fn, a built-in or a value) whose first parameter is a model, and which has no mut parameters (E1105). The other arguments are checked as for a call f(m, args…); m must be the model of f’s first parameter (E0904). f must return a [] of floats, such as a loss (E1106), and the value has that type. f may take a mode, as a loss with dropout does: value_and_grad(loss, m, x, y, Train(seed=step)); gradients go through the values dropout keeps.

Which params f reads is worked out from its code: m.w1 reads w1, and passing m to another fn reads what that one reads, also through pipes and nested value_and_grad. Passing m anywhere else, such as to Python, may read any param. If f reads no param of m, its gradient would always be zero, which is E1107; each param it never reads gets the warning W1101.

Gradient types

Grad[M] is the type of a gradient with respect to the model M. It has a field for each param of M, with the param’s type, read as g.w1. Buffers aren’t differentiated, so it has no fields for them (E0901). Grad[M] can be a parameter or return type; a Grad of something other than a model is E0303, and so is as Grad[M].

Gradients of the same model can be added and subtracted, -g negates one, and a gradient can be multiplied by a [] or a number on either side, or divided by one. norm(g) is the norm of all the gradient’s values together, a [] NonNeg. Anything else, such as g * g, g + x for a tensor x, 2 / g or mean(g), is E1108; adding gradients of different models is E0904.

Gradients at run time

value_and_grad compiles to rt.value_and_grad(f, m, args…), which runs f with autograd enabled, even inside torch.no_grad(), and calls torch.autograd.grad with the model’s params. Nothing reads or writes .grad, so there is nothing to zero, and no train() or eval() to switch; randomness comes from a mode argument, which value_and_grad passes on to f like any other. The value is detached. A gradient is an rt.Grad: its fields are attributes and items (g.w1, g["w1"]), in declaration order, and it supports the arithmetic above. Each gradient tensor has its param’s shape and dtype, and is where the param is, the declared device; a param that f doesn’t use gets zeros. A Grad[M] parameter checks its argument is a gradient of M with the params’ shapes and dtypes (Runtime checks).