19: Python calls and options on several devices

E1901

Python call given data split across the devices. With several devices, Python calls run once, in the controller, as they would on one device. A batch is split across the devices, each holding some of the examples, and so is every value with its batch axis; giving one to a Python function would gather the batch through the controller’s CPU at every step.

devices cuda[0 1 2]

import python augment

dim D = 4

data Ex:
    x: [D]

fn augmented(batch: [b] Ex) -> [b D]:
    return augment.random_crop(batch.x) as [b D]

Move the call into the dataset, so it runs on each example as it’s loaded, or declare the function with extern fn, as extern fn random_crop(x: [b D]) -> [b D] = "augment.random_crop", so each device runs it on its own examples. See Python calls on several devices.

E1902

Split axis whose size the devices don’t divide. split H gives every device the same share of each axis of size H, so the number of devices must divide H.

devices cuda[0 1 2 3 4]
    split H

dim H = 12

Use a number of devices that divides it, here 2, 3, 4, 6 or 12, or split another dim. See Options for several devices.

E1903

devices option given twice. Each option is given once. The devices are in a row, so they can split values along one axis only: a second split, of another dim, is an error too, since a value with both axes couldn’t be split along both.

devices cuda[0 1]
    split H
    split D

dim D = 8
dim H = 2

Keep one split. See Options for several devices.

E1904

split of something other than a dim. split NAME splits every axis whose size is the dim NAME, so NAME must be a dim the file declares: an axis name, such as a generic b, isn’t one.

devices cuda[0 1]
    split heads

fn f(x: [heads n]) -> [heads n] = x * 2

Declare the size as a dim, as dim H = 12, and split that. See Options for several devices.