Modes and randomness

Modes

Mode is the type of a mode: Train(seed=s), for training, whose random operations draw from the seed s, or Eval, for evaluation, whose random operations do nothing. Random operations, of which dropout is the first, take a mode as an argument. There is no global training flag and no global random state, so a function that draws random numbers says so with a Mode parameter, and every call to it passes one:

fn forward(m: Mlp, x: [b D], mode: Mode) -> [b C]:
    return x @ m.w1 |> gelu |> dropout(0.1, mode) |> matmul(m.w2)

fn train_step(m: Mlp, x: [b D], step: int) -> [b C] = forward(m, x, Train(seed=step))

fn evaluate(m: Mlp, x: [b D]) -> [b C] = forward(m, x, Eval)

Mode is a parameter or return type, and a mode can be bound to a name, passed and returned like any value. A call that leaves out a Mode argument is E1201, and anything other than a mode where one is expected is E1202; a Python value fits, and is checked when the function runs (Runtime checks). A mode isn’t a tensor (E0412), and Mode in another type, a Mode parameter with a default, Train without (seed=…), and Eval(...) are E0303 or E0306. seed is an int, given by name (E0406; E0502 or E0805 for something other than an int). Mode, Train and Eval are built-in names, which a declaration or a value of the same name hides.

From Python, the modes are shapescript.Train(seed=…) and shapescript.Eval, both instances of shapescript.Mode. They are immutable.

dropout

dropout(x, p, mode) takes a tensor of any shape, a float p and a mode, and gives a float tensor of x’s shape and dtype whose domain is Real. With Eval it returns x itself. With Train(seed=…) each value is zeroed with probability p and the others are divided by 1 - p, so the mean is unchanged; gradients go through the values kept. p is at least 0 and less than 1, or dropout raises ValueError when it runs.

batch_norm

batch_norm(x, mode, over=C, weight=w, bias=b, running_mean=mu, running_var=var, eps=1e-5) is batch normalization, a stub in nn.ss: each index of the axis over, the channel, is normalized over every other axis of x, as torch’s BatchNorm1d and BatchNorm2d normalize over every axis but the channels. weight, bias, running_mean and running_var are [C] and may be left out. The mode decides which statistics it uses: with Train(seed=…), the batch’s mean and biased variance; with Eval, running_mean and running_var, which must be given then (ValueError otherwise). It draws nothing at random, and keeps no state: the running statistics are values like any other, such as a model’s buffers, and a training step updates them itself, with mean over the batch:

fn track(mut m: Net, x: [b D]):
    h = x @ m.w1
    mu = mean(h, over=b)
    m.var = 0.9 * m.var + 0.1 * mean(square(h - mu), over=b)
    m.mean = 0.9 * m.mean + 0.1 * mu

On several devices, its statistics with Train are the whole batch’s.

Random draws

The same seed gives the same random numbers, and two random operations never share a draw, even with the same mode:

• Paths. A training mode carries a path as well as its seed. Every call that passes a mode on, to a ShapeScript function, a prelude stub such as dropout, or a Python function, extends the path with the callee’s name and how many times this call of the enclosing function has already passed a mode to it: forward#0, then forward#1; dropout#0, then dropout#1. value_and_grad(f, m, …, mode) passes the mode on as the call f(m, …, mode) would, extended with f’s name, so a training step’s loss draws the masks that calling the loss function there would. Inside a loop, the count is per iteration, and the name says which iteration: forward@4#0 is the first call of forward in the iteration of for epoch in a..b: whose epoch is 4, and forward@4.17#0 the first in the 17th (from 0) iteration of a loop inside it over anything else, such as batches. So the two dropouts in one function, two calls of forward with the same mode, and each iteration of a loop that calls forward draw different masks, and a loop’s draws depend on its iterations, not on where it started: a loop over 3..5 draws what one over 0..5 draws in its last two iterations. A mode the checker sees only as a Python value isn’t extended, so give it its type, with a Mode parameter, before using it twice. A notebook cell counts its own calls outside any fn. Calling the same function from Python twice with the same mode gives the same masks both times.
• Generators. Each draw has a torch.Generator of its own on the declared device, seeded with the first 8 bytes of the SHA-256 of "{seed}/{path}", the path’s parts joined with /, as a model field’s seed is made (Compiled models). Nothing reads or changes PyTorch’s global random state. Masks are drawn as f32 whatever the input’s dtype, on the declared device, so a seed gives the same mask on the same kind of device; CPU, CUDA and MPS generators give different numbers.

Paths count calls, not source positions, so editing other functions, or adding a line before a call, doesn’t change its draws. Adding a call to the same function before it, in the same iteration, does.