Debugging

ShapeScript is debugged with Python’s debugger, debugpy, which VS Code’s Python Debugger extension and Jupyter use. The code the import hook, python -m shapescript and the cell magic run reports the .ss file (or the cell) and its lines, so a breakpoint in a .ss file stops there, the call stack mixes .ss and Python frames, and stepping goes from one language to the other: step into a ShapeScript line that calls Python and you’re in the Python function; step out and you’re back on the line; step into a call from Python and you’re on the fn line of the ShapeScript function.

Every line is a place to stop. A pipe written over several lines stops on the line before its stages and on each stage. The fn line is where the argument checks run.

When Python imports a .ss file that has errors, through the import hook or the module generated beside it, the ShapeScriptError it raises, which shows every diagnostic, is raised from the file’s first error: the last entry of its traceback is that line of the .ss file, as File "model.ss", line 5, in <module>, with the error’s span underlined from Python 3.11 on. So a debugger that stops on uncaught exceptions, as VS Code’s does by default, stops on the broken line of the .ss file and shows the error there, not in the Python that imported it.

pip install shapescript installs a debugpy extension (pydevd_plugins/extensions/pydevd_plugin_shapescript.py), which debugpy loads whatever starts it. It:

• hides the names the compiler makes up, which all start with _ss_ (temporaries, sizes, the device, the runtime module);
• shows a tensor in a .ss frame with its sizes, named after the axes of its declared type, before its values: x = [b=2 D=4] float32 tensor([[1., …, and gives each parameter and value its ShapeScript type, as hover does ([b D] Prob over D), which VS Code shows when you point at the variable, and so each field of a model it holds (m.w);
• does the same for functions that share a name (overloaded functions), whose Python names are _ss_forward__Mlp and the like, and for each instance of a generic function, _ss_fit__Mlp__Digit, whose values have the types of that instance; an instance’s code is shown at the call it was made for, as its lines are. Imports of other .ss files are found as the import hook finds them;
• steps over the shapescript package, the runtime that generated code calls, even with "justMyCode": false or an editable install. A launch configuration’s rules come first, so {"module": "shapescript", "include": true} steps into it.

python -m shapescript program.py [args…] runs a Python program with the import hook installed, so it needs no shapescript.install(); python -m shapescript model.ss [args…] imports the file and calls its fn main(), which takes no parameters; python -m shapescript -m module runs a module. VS Code’s shapescript launch configurations run these under debugpy (see the extension’s README); any other debugpy configuration works too, such as the Python extension’s “Python File”, if the program installs the hook itself.

In a notebook, the %%shapescript cell magic notices when Jupyter’s debugger (ipykernel’s, which VS Code and JupyterLab use) has written the cell to a file to set breakpoints in it, and its code then reports that file and the cell’s lines, counting the %%shapescript line, so breakpoints in the cell stop. Types come from the cell alone, so values whose types use what earlier cells declared show only their sizes.

Not debuggable yet: the worker processes of several devices, where each device’s work runs (debugpy debugs only the controller’s process, so breakpoints in fns that run on the workers don’t stop); and Python written by shapescript build, which is plain Python whose lines are its own.

The tensor view

In VS Code, a stopped program’s tensors are drawn in the Tensor view, a section of the Run and Debug view beside Variables, Watch and Call Stack, which shows while debugging. View Tensor adds an expression to it and shows it: from the context menu of a variable in the Variables or Watch view, from the editor’s context menu on a name (or a selection) while stopped, or from the Command Palette or the + in the view’s title bar, which ask for an expression. The view keeps its expressions as Watch does, one tab each, and draws the selected one; its expression box changes it. The expression is evaluated in the selected frame, as a watch is, so it can be any Python expression of the frame’s names, such as m.w or att[0], and a NumPy array works too. When the program stops again, or another frame is selected in the call stack, the view shows the expression there; while the program runs, it keeps what it drew, marked as the values of the last stop. In the side bar it’s one column, the heatmap, then the checks and statistics; made wider, or moved to the panel, it puts them side by side.

The tensor is drawn as a heatmap of two of its axes, zoomed with the mouse wheel and panned by dragging, from the whole tensor down to single elements, which show their values as numbers once the cells are big enough. Pointing at a cell shows where it is, by axis name (att[b=0, H=2, n=20, m=41]), and its exact value, written with the fewest digits that are the same number in the tensor’s dtype; a cell that covers several elements shows their mean, least and greatest. Each axis has a role, chosen in the view: the heatmap’s rows or columns, small multiples (a heatmap for each of its indices), color channels (for an axis of 3 or 4), a slider, or a reduction over it (mean, max, min, sum, largest |x|, norm or std). A group axis whose parts’ sizes are known, such as (gh gw), can be split into its parts, which then have roles of their own. A tensor of one axis is drawn as a line, with the least and greatest values under each pixel as a band, and a scalar as its value.

The tensor’s ShapeScript type, from the checker, as hover shows it, picks the layout and the colors. Values of other Python code, which have no type, are drawn from their dtype alone, with their axes numbered.

• An axis made of two named sizes followed by one more axis, as the tokens of a grid of patches [b (gh gw) D], is split, and its parts are the rows and columns, each cell the token’s norm over D.
• Otherwise, where the checker knows a matrix is a background outside a band of diagonals, as a causal mask causal[n m] is false above its diagonal, and what masked_fill and softmax make of one is -inf or 0 there (its structure, which hover draws), its rows and columns are the heatmap’s.
• Otherwise Logit, Prob or LogProb over m puts m on the columns and the axis before it on the rows.
• Otherwise an axis of size 3 before the last two is the color channels of an image whose values are Unit or u8, and a slider for other values; the last two axes are the rows and columns.
• Otherwise the last two axes are the rows and columns.

Then the nearest remaining axis before the rows, if it has 2 to 16 elements, is small multiples, as H is in [b H n m], and the rest are sliders at 0. Colors: a mask (Mask, Binary) has two; Index[V] with at most 8 values has the categorical palette’s colors; values that can’t be negative (Prob, Unit, NonNeg, LogProb) have a sequential scale from the color nearest the background, so small values recede; and others have a diverging scale around 0, blue below and red above. The range is the domain’s when it’s bounded, as [0, 1] for Prob; otherwise the values’ 0.5th to 99.5th percentiles while each cell is an element, and the shown cells’ range when cells are reductions or cover several elements. The histogram beside the heatmap is its legend: dragging its ends sets the range, and double-clicking resets it. A cell that covers several elements shows, by default, the greatest of values that can’t be negative, the value of largest |x| (with its sign) of others, and the mean of masks and indices, so that a single outlier shows at any zoom; the view’s Cells menu chooses the mean, least or greatest instead.

The checks list what the type promises and whether the values keep it, with Show, which goes to the first or worst element that doesn’t:

• no NaN or ±∞ (-inf where a mask’s background is -inf, or in Additive values, isn’t counted);
• every value is in the domain: in its range, as [0, 1] for Unit or Prob, at most 0 for LogProb, 0 or 1 for Binary, 0 or -inf for Additive, and whole numbers in [0, V) for Index[V];
• each row of Prob over m sums to 1 along m (each row of LogProb’s exponentials does), to within 1e-3 or twice the dtype’s epsilon, whichever is larger;
• the background its structure promises is there: after a causal mask, att[b, n, m] is 0 wherever m > n.

On the heatmap, NaN and ±∞ are hatched in red, values outside the domain, and those that aren’t the background their structure promises, in orange, and that background is hatched in gray, with its edge dashed; rows of a Prob that don’t sum to 1 are marked at the heatmap’s left edge. Where cells are smaller than a few pixels, a ring marks each flagged one, so one NaN in a billion values shows. Beside the checks are the values’ statistics, and, for floating-point dtypes, how far the largest |x| is from the dtype’s largest number (its headroom, in bits) and how many values are subnormal.

The tensor never leaves the device it’s on. View Tensor evaluates shapescript._tensorview.open(...) in the frame, which keeps a reference to the value and starts a server in the program, on 127.0.0.1 with a random token; the view asks it for a summary and for tiles of 256 × 256 cells, which it computes where the tensor is, reading it in blocks whose temporaries take at most an eighth of the GPU’s free memory, and sends only those. A cell covers 2ᵏ × 2ʲ elements, so a few tiles cover the whole tensor when zoomed out, and the cells keep each element’s NaN and ±∞ and whether it’s in the domain. On a 5090, a summary of 8 GiB of bf16 takes under a second, and a zoomed-out tile of all of it about 160 ms. The server’s threads are marked as the debugger’s own, so debugpy neither traces nor suspends them, and on CUDA they read on a stream of their own after the work the stopped thread had queued, so the program’s streams are untouched. A value that’s being recorded into a CUDA graph has no values yet, and isn’t drawn. Random samples, for the percentiles, use a generator of their own, so the program’s random numbers don’t change. When the program goes on, the view releases the value; when the debugger goes away, the server releases those it read.

Not yet: values on another machine than VS Code’s extension (with Remote SSH and containers, the extension runs where the program does), complex, sparse and quantized tensors, values inside a map (which PyTorch’s vmap sees as batched tensors), and the worker processes of several devices.