Changing models with mut

mut parameters

A function may change a model, or optimizer state, only if it is a mut parameter, and every call writes mut before that argument, so both the signature and the call show what changes:

fn step(mut m: Mlp, x: [b D], y: [b] Index[C], lr: float):
    l, g = value_and_grad(loss, m, x, y)
    m.w1 = m.w1 - lr * g.w1
    m.w2 = m.w2 - lr * g.w2

fn fit(mut m: Mlp, x: [b D], y: [b] Index[C], steps: int):
    for i in 0..steps:
        step(mut m, x, y, 0.1)

Only a parameter whose type is a model or optimizer state, such as AdamW[Mlp], can be mut; a tensor, host scalar or schedule is a value, and a function returns a new one instead (E1101). An extern fn’s parameters can’t be mut (E1101). At a call, an argument for a mut parameter needs mut (E1103), and mut before an argument whose parameter isn’t mut, or before an argument of a built-in or of Mlp.init, is E1101. Only a name can be passed as mut (E1101), so the input of a pipe can’t be. A Python function can change anything, so an argument of one may be marked mut; what it passes must be something the caller may change.

What counts as a change

A model is changed by assigning one of its fields, m.w1 = e, or by passing it as mut, as to update; optimizer state only by passing it as mut. Reading fields, and passing a model to a parameter without mut, don’t change it. e must have the field’s shape and dtype (E0403, E0401, E0502), and fit a buffer’s domain (E0501); only a model’s fields can be assigned (E0303), and only declared ones (E0901).

A function may change:

• its mut parameters;
• models and optimizer state it makes or gets from calls, such as Mlp.init(seed=0) or adamw(m, lr=3e-4), unless they may be a parameter without mut (below);
• in a notebook cell, outside any fn, any value.

Anything else is E1102: a parameter without mut, and a value that may be one, because it was assigned from one (a = m) or computed from one with the same type (a function can return its argument). After an if or a loop, a value that may be such a parameter on some path is treated as one. A Python variable seen inside a function can’t be changed either.

A mut parameter can’t be given a new value, as in m = Mlp.init(seed=1) or by a for, since the caller would still have the model or state it passed (E1104).

Changes at run time

A field assignment compiles to rt.set_field(m, "w1", e), which copies the new value into the field’s tensor in place, under torch.no_grad(). So from Python, a mut argument is changed in place: the caller’s model object changes, its params stay the same nn.Parameter objects, and optimizers and state_dict() still refer to them. The new value isn’t part of any gradient computation.