17: Export

E1701

Python in an exported function. An export fn, and every function it calls, can’t use a name from import python, or anything reached through one: a call such as torch.load("w.pt"), a value such as torch.pi, or a pipe stage such as |> torch.tanh. In a notebook, a Python variable is the same. torch.export records torch operations only, and a Python call could do anything.

dim D = 768

import python torch

export fn load_qkv() -> [D 3*D]:
    return torch.load("w_qkv.pt") as [D 3*D]

Pass the value in as a parameter, or give a torch function a type with extern fn, as in extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc". See Export functions.

E1702

extern fn outside torch, called by an exported function. An exported function can call an extern fn whose Python function is in torch or one of its submodules, or in math. One from another module could run any Python, which torch.export can’t record.

extern fn mystery(x: [...]) -> [...] = "mylib.mystery"

export fn f(x: [n]) -> [n]:
    return mystery(x)

See Export functions.

E1703

as on a Python value in an exported function. In an exported function every value must come from torch operations. A Python value, such as one an earlier line got from Python (itself E1701), can’t be given a type with as there.

dim D = 768

import python numpy

export fn scale(x: [n D]) -> [n D]:
    w = numpy.ones(768)
    return x * (w as [D])

Pass the value in as a parameter of the function. See Export functions.

E1704

Value or operation with no torch-only lowering in an exported function. Models, modes, logs records, sinks, gradients, optimizer state, schedules, datasets and batches exist only on the host, and so do the built-ins that make or use them, such as print, save, value_and_grad and update. Changing a model’s field and a for loop over anything but a range are the same.

dim D = 64

export fn f(x: [n D]) -> [n D]:
    print(x)
    return x

An exported program is for inference: pass the model’s fields in as tensors. See Export functions.

E1705

Parameter of an export fn that isn’t a tensor with a shape. The parameters of an exported function are the program’s inputs, and torch.export traces it with example tensors whose shapes, and which of whose axes may change, come from the parameters’ types. A parameter without a shape, a model, or an int, float or bool without a default has none.

export fn f(x: [n], k: int) -> [n]:
    return x * k

Give it a tensor type such as [b n D], or a host scalar a default, which the exported program keeps. See Export functions.

E1706

return inside an if on a tensor in an exported function. An if on a [] Mask becomes torch.cond, whose branches each give values, which come out of the if; a branch can’t return from the function.

export fn f(x: [n]) -> [n]:
    if sum(x) > 0:
        return x
    return -x

Assign the value in both branches and return it after the if: if sum(x) > 0: y = x / else: y = -x, then return y. See Control flow in exported programs.

Keyboard shortcuts