Notebooks

ShapeScript works in Jupyter and VS Code notebooks through a cell magic. Exploration is where shape bugs are born, so having the checker there helps most.

Setting up

A notebook runs its cells in a kernel, a Python process, and for ShapeScript cells it must be a Python with ShapeScript installed: your project’s environment. Chapter 1 sets it up (Notebooks): pip install ipykernel in the environment, then, in VS Code, Select Kernel at the top right of the notebook › Python Environments › .venv. In Jupyter in the browser, install jupyterlab into the environment and start jupyter lab from it.

To check which Python a notebook runs, run import sys; sys.executable in a cell: it should be inside your .venv. If %load_ext shapescript says No module named 'shapescript', it isn’t; see Appendix D.

Loading the magic

In a Python cell:

%load_ext shapescript

From then on, a cell that starts with %%shapescript is ShapeScript. A notebook declares its devices like a program, once:

%%shapescript
devices cpu
dim D = 768
fn mlp(x: [b D], w1: [D 4*D], w2: [4*D D]) -> [b D] = x @ w1 |> gelu |> matmul(w2)

Cells share what they declare

Unlike a file, a cell may also hold statements outside any function, and they run in order. Each cell is checked together with the cells that ran before it, so D and mlp from the cell above are known here:

%%shapescript
import python torch
x = torch.randn(8, 768) as [8 D]
y = mlp(x, torch.randn(768, 3072) as [D 4*D], torch.randn(3072, 768) as [4*D D])
y

The last line is a tensor, so the cell shows its shape diagram next to its value. Python variables of the notebook are visible too, as Python values, so a tensor made in a Python cell comes in with an as:

%%shapescript
z = mlp(batch as [4 D], torch.randn(768, 3072) as [D 4*D], torch.randn(3072, 768) as [4*D D])

(Here batch was made in a Python cell, batch = torch.randn(4, 768).) And ShapeScript values are Python variables afterwards, so a Python cell can plot y or save it.

A cell with errors runs nothing

A cell with a type error shows the error as its output, and runs nothing, so a half-run cell never leaves the notebook in a confusing state:

%%shapescript
w = torch.randn(768, 10) as [D 10]
bad = mlp(y @ w, w, w)

The output names [1], [2] and so on for the cells, in the order they ran, where a file would have a file name.

Running a cell again

Running a cell again replaces what it defined. If you change dim D = 768 to dim D = 512 and run that cell, later cells check against the new D the next time they run. Values made with the old D, such as y, can’t be used until the cell that made them runs again: their types no longer match the code, which is error E0701, rather than a silent mismatch.

Devices are the exception: a notebook’s devices can’t change once declared. Restart the kernel to change them.

Debugging in notebooks

Jupyter’s debugger, which VS Code and JupyterLab use, stops at breakpoints in %%shapescript cells, and steps into the Python they call, as it does for .ss files (Chapter 13).

Exercises

1. In a notebook where one cell declared dim D = 768, write a cell that declares fn norm(x: [b D]) -> [b D], a layer norm over D, and a second cell that calls it on a Python tensor t of shape (2, 768).
2. What happens if you run the second cell before the first?