ShapeScript for VS Code

Highlighting for .ss files, plus what the ShapeScript language server (shapescript-lsp) provides:

• errors and warnings as you type;
• hover: the type of any expression or pipe stage with a diagram of its shape, drawn in your theme's text color, a dim's value and ## docs, and where an axis comes from;
• hover on an activation, built in or of your own (an elementwise function, or any function of one tensor and host scalars that only does elementwise maths to the tensor, such as fn act(x: [b D]) -> [b D] = x * sigmoid(1.702 * x)): its formula, its curve with the derivative dashed, and the range of its values, using the numbers written in the call, as in leaky_relu(0.2);
• hover on a function, of your own or built in, or on a call of a built-in that computes, such as argmax or cross_entropy, |> into one, or to between domains: a picture of what it does, run on a small example in your Python interpreter (see Pictures), under a call's types before and after, with what changed in color;
• hover on a built-in that only moves values, such as rearrange or split, -> in its pattern, or |> into one or into a function of your own: the tensors before and the result after, side by side in one image with an arrow across the gap between them, and what the step changed in color (see Before and after), after a built-in's signature, what it does and an example;
• hover on a dim: where it's used, grouped by function, with each value's type;
• go to definition, find references, and rename (refused if it would change what a name refers to); go to definition on a built-in opens its declaration read-only (see below);
• completion as you type, everywhere a name goes, ranked by the type the place wants, with calls as snippets, and signature help in calls (see Completion);
• completion after m., for a model: its fields, and the functions that take it first, from the file, its imports and the workspace's other .ss files, which become calls (see Completion after m.);
• go to definition between ShapeScript and Python, both ways (see Python and ShapeScript);
• the shape at the end of each pipe line, as an inlay hint, whose hover shows the diagram;
• different colors for dims, axes, domains and values;
• math in comments: hover shows it typeset, in your theme's text color, and typing \ in $...$ offers the LaTeX commands;
• debugging: breakpoints in .ss files, and stepping between ShapeScript and Python (see Debugging);
• the tensor view: View Tensor on a variable of a stopped program draws it in the Run and Debug view as a heatmap you zoom from the whole tensor down to its numbers, laid out and checked by its ShapeScript type (see The tensor view).

Setup

The language server comes inside the extension, so .ss files get errors, hovers and the rest as soon as it's installed. Running and debugging ShapeScript, the generated Python modules and notebooks also need the Python package in your project's environment:

python3 -m venv .venv && source .venv/bin/activate    # Windows: .venv\Scripts\activate
pip install torch numpy
pip install shapescript

Then select that environment with Python: Select Interpreter: the server finds installed ShapeScript packages and Python sources through it, and the debugger and notebooks run in it. The Getting Started chapter of the ShapeScript book goes through every step.

Debugging

ShapeScript runs as Python, and is debugged with Python's debugger, debugpy, as Python is. Breakpoints can be set in .ss files; the call stack shows .ss and Python frames together; stepping into a ShapeScript line that calls Python goes into the Python function, stepping out comes back to the line, and stepping into a ShapeScript function from Python stops on its fn line. Each line of a pipe written over several lines is a step. The variables view shows ShapeScript's names with their types and sizes (x = [b=2 D=4] float32 tensor(…), type [b D]), without the _ss_ names the compiler makes up, and stepping skips ShapeScript's runtime. That comes from a debugpy extension that the shapescript Python package installs, so it works whichever way debugpy starts.

To debug, either:

• use the extension's ShapeScript debug type: ShapeScript: Debug File in the editor's run menu, or F5 with a .ss or .py file open and no launch.json, or a configuration such as

  {
    "type": "shapescript",
    "request": "launch",
    "name": "Train",
    "program": "${workspaceFolder}/train.py",   // or a .ss file with `fn main()`, or "module": "..."
    "args": []
  }
  

  It runs python -m shapescript program args…, which installs the import hook first, so a Python program can import model from model.ss without calling shapescript.install(), and a .ss file's fn main() runs. With the Python Debugger extension installed, the session is handed to it, with its own debugpy and the interpreter selected in the Python extension. Without it, the extension starts debugpy itself, with the configuration's python, the shapescript.debug.python setting, or python3 on your PATH, which needs pip install debugpy shapescript;

• or use any Python debug configuration ("type": "debugpy") for a program that calls shapescript.install(): VS Code sends breakpoints in .ss files to every debug session, and debugpy finds the code that reports them.

When the program imports a .ss file that has errors, the debugger stops on the file's first error, with every diagnostic in the exception's popup, as it stops on any uncaught exception (Uncaught Exceptions in the Breakpoints view, on by default).

In a notebook, breakpoints in a %%shapescript cell stop when the cell is debugged (Debug Cell), as in a Python cell. This is tested against ipykernel's debugger, which VS Code's notebooks use, through the protocol rather than in VS Code itself.

Not yet: debugging the worker processes of a program on several devices (only the controller's process is debugged), and stepping into Python written by shapescript build, which has its own lines.

The tensor view

The Tensor view is a section of the Run and Debug view, beside Variables, Watch and Call Stack, while debugging. While the program is stopped, View Tensor adds a tensor to it and shows it: right-click a variable in the Variables or Watch view, or a name in the editor (or select an expression), or click + in the view's title bar, or run ShapeScript: View Tensor, and type an expression of the selected frame, such as att or m.w. NumPy arrays work too. The view keeps a tab for each expression, as Watch keeps its entries, and follows the selected one from stop to stop and from frame to frame, keeping its picture while the program runs. In the side bar it's one column; widen the side bar, or drag the view to the panel, for the heatmap and the checks side by side.

It draws the tensor as a heatmap of two of its axes: scroll to zoom, drag to pan, and zoom in until the cells show their numbers; point at a cell for its index by axis name and its exact value. Each axis has a role (rows, columns, small multiples, color channels, a slider, or a reduction over it), which the tensor's ShapeScript type chooses to begin with: [b H n m] Prob over m gets a heatmap of n by m per head H, with b as a slider; an image [b 3 h w] Unit is drawn in color; patch tokens [b (gh gw) D] are drawn as their grid, each token's norm over D. The type's domain chooses the colors, and the checks beside the heatmap say whether the values keep what the type promises: no NaN or ±∞, values in the domain, rows of Prob that sum to 1, zeros where a causal mask was applied. Show goes to the first or worst value that doesn't. The histogram is the colors' legend; drag its ends to change their range. Keys, with the heatmap focused: + and - zoom, 0 fits, the arrows pan, and [ and ] move the first slider.

The program computes everything the view shows on the device where the tensor is, in tiles of 256 × 256 cells, and sends only those, so tensors of many gigabytes on a GPU are drawn as quickly as small ones. It does this in threads the debugger leaves alone, on a stream of its own, without changing the program's random numbers. The program has to run on the machine the extension runs on (with Remote SSH or Dev Containers, it does). See The tensor view in the reference for the details.

• a Python module beside each .ss file, so Python code that imports it type-checks (below).

Python modules

Pylance doesn't know about .ss files, so in a Python script from model import f is Import "model" could not be resolved, even though it runs once shapescript.install() has. So when you open or save model.ss and it has no errors, the language server writes model.py beside it: typed declarations of its dims, models and functions, whose docstrings show the ShapeScript signatures, so Pylance resolves the import, completes the names and shows the shapes on hover. When Python imports model.py, its last line loads the compiled model.ss instead, so import model works without shapescript.install() too. See Python modules for .ss files.

• The first line of model.py says it's generated. A model.py without it is yours and is never replaced; the server shows warning W2501 on model.ss instead.
• shapescript.pythonModules (on by default) turns this off; it takes effect when the server restarts (Developer: Reload Window).
• The extension turns on explorer file nesting (explorer.fileNesting.enabled) and adds the pattern "*.ss": "${capture}.py", so model.py is shown inside model.ss. VS Code adds an extension's patterns to its own and to yours, key by key, so none of them is lost. To keep nesting off, set "explorer.fileNesting.enabled": false in your settings.
• The modules are derived files: .gitignore them, or commit them so the imports type-check without an editor, and check them in CI with shapescript pymodules --check.

Completion

Suggestions show as you type in .ss files, as for any language (editor.quickSuggestions), and also inside a call's tab stops once you've accepted a call: the extension sets editor.suggest.snippetsPreventQuickSuggestions to false for ShapeScript, as the Go extension does for Go. The list has the values in scope with their types, the file's declarations, the built-ins, keywords where they go, and once you type a letter, the names of the workspace's other .ss files with the import they need; the right side says where an item comes from (metrics · adds its import). What fits the place's type comes first, and the best one is selected:

• in a call's argument, values of the parameter's type, and the keyword arguments not given yet (over=, then the operand's axes);
• after |>, the functions that take the value piped in, as stages;
• in a pattern, the input's axes; in a type, sizes, axes, domains and dtypes;
• after ., a batch's fields, an imported file's declarations, a model's (see below), Mlp.init, and a Python module's names, read from your interpreter (your own modules are read, not run);
• after import and import python , the files and modules there are.

The list opens by itself after ., [ in a type, ( and , in a call whose argument's type is known, = after a keyword argument's name, and a space after |> or import. Signature help shows the signature of the call you're typing, with the current parameter highlighted; all overloads are listed. See Completion and signature help.

Completion after m.

ShapeScript has no methods: a model's functions are ordinary functions that take it first, and a call picks the one for its model. So after m., where m is a model, the list has the model's fields and those functions, from the file, the files it imports and every other .ss file in the workspace (except files that import this one). Choosing forward replaces m. with forward(m, and the other arguments as tab stops; a function of another file is written attention.forward(m, …), and if the file doesn't import attention yet, import attention is added after its imports. The description on the right says which file a function comes from and whether it adds an import. A generic function (fn fit(mut m: M, …)) is listed too, marked "needs loss for Mlp" if the model lacks a function it calls.

Python and ShapeScript

• From ShapeScript: Ctrl-click on the module of an import python, on torch.load or another member of an imported module, or on an extern fn's path string opens the Python source. The language server finds it with the interpreter the Python extension has selected for the workspace, if it's installed, and follows changes to it; otherwise with shapescript.debug.python, else python3. Finding a name in a module doesn't run the module, unless the module only imports it (as torch does load).
• From Python: Ctrl-click on a name imported from a .ss module, as causal_attention in from attention import causal_attention (or attention.causal_attention after import attention), opens its declaration in attention.ss, and so does Ctrl-click on a declaration in the generated attention.py. The generated module never shows on the way: Pylance answers with it, so the extension adds no second answer (which would make a peek list), and shows attention.ss at the same declaration in place of the generated module's editor; Back skips it. A generated module you open yourself stays open. Peek Definition (Alt+F12) still shows Pylance's answer. For this the extension activates in workspaces that contain .ss files.
• Python's uses of ShapeScript: the references of a function, model, dim or record of a .ss file (Ctrl-click on its declaration, or Find All References) include the Python code that uses it, which Pylance finds; without a Python language server, the extension reads the workspace's Python files that import the module.

Running it from source

Open the repository in VS Code and press F5. That builds the language server (cargo build -p shapescript-lsp) and the extension, then opens examples/ in a new window with the extension loaded.

The extension looks for the server in this order:

1. the shapescript.server.path setting;
2. server/shapescript-lsp inside the extension (for published builds);
3. target/debug or target/release in this repository, when the extension runs from source;
4. shapescript-lsp on your PATH.

Tests

npm ci
npm run typecheck
npm test              # the grammar, tokenized the way VS Code does it, the math typesetter, and Python imports
npm run test:e2e      # the extension in a real VS Code window
SHAPESCRIPT_VSIX=shapescript-linux-x64.vsix npm run test:e2e   # a packaged extension, with its bundled server

npm run test:pylance checks the Python modules (PY-4) and going between Python and ShapeScript with the real Python extension and Pylance, which it installs from the Marketplace into a temporary extensions folder of the test VS Code (never your own), on a copy of tests/python/playground with the repository's .venv (or SHAPESCRIPT_E2E_PYTHON) as the interpreter. It needs cargo build -p shapescript-lsp -p shapescript-cli first. SHAPESCRIPT_E2E_ONLY=<regex> runs only the end-to-end tests whose names match. When the test window isn't in front (macOS doesn't let it take the focus from another application, and VS Code on Linux has no command that asks for it), typing, accepting a suggestion and Ctrl-click, which VS Code does only in the focused editor, are done through the same providers instead: inserting the text, applying the item the list would select, and opening the definition.

test:e2e needs cargo build -p shapescript-lsp first, and a display (xvfb-run -a npm run test:e2e on a machine without one). Its debugging test runs a program in tests/debug under debugpy with a Python that has debugpy and shapescript installed: SHAPESCRIPT_E2E_PYTHON, else the repository's .venv, else it's skipped. The Python extension isn't installed in the test's VS Code, so the test covers the extension starting debugpy itself, not handing the session to the Python Debugger extension. It downloads VS Code unless VSCODE_EXECUTABLE points at an installed one, such as /usr/share/code/code. Unset ELECTRON_RUN_AS_NODE if your shell sets it, or VS Code starts as plain Node.js. The editor gets a fresh profile and a temporary workspace, so it doesn't touch your settings or files.

The grammar's lists of built-in functions and dtypes copy the compiler's. crates/lsp/tests/grammar.rs fails if they drift apart.

Math

VS Code doesn't typeset math in hovers, so the extension asks the server for math as LaTeX (initializationOptions: { math: "client" }) and typesets it with MathJax 3, bundled into dist/extension.js. MathJax draws each formula as an SVG image whose glyphs are paths, in the hover's text color, and the extension keeps the images it has drawn, by formula and color. Inline math is sized so its x-height matches the hover's text (7 pixels per ex); a formula on lines of its own, larger (8.5), since its fractions and subscripts are smaller than its text. MathJax makes the minified extension about 2.1 MB, where it would be 0.44 MB without it.

Declarations of built-ins

Go to definition on a built-in, such as Ctrl-click on gelu or rearrange, opens its declaration: prelude.ss and nn.ss, the prelude's stubs, or builtins.ss, which the server generates with the documentation its hovers show. The extension passes its global storage folder to the server as libraryDir; the server writes the three files there, read-only, and answers with their file:// locations. The extension's definition middleware turns those into shapescript-library:/builtins.ss and the like, which a content provider serves from the folder, so they open as read-only documents. Other editors get the files themselves, in a cache folder unless they send libraryDir (see the language reference).

Pictures

A function's hover, and a call's, has a picture of what it does (HOV-4): the language server (crates/lsp/src/picture.rs) makes the function's file small, with its dims made small, compiles it to standalone Python, and runs the function on an example made from its parameters' types in a worker of your Python interpreter (crates/lsp/src/runner.rs), which needs PyTorch. It draws the values going in and coming out when the result's type says something about them, and otherwise which values of the result change when one value going in does; the language reference has the rules. The picture is an SVG in currentColor and the accent orange, which the extension recolors as it does diagrams.

Running a function is code execution, so the extension tells the server to do it (initializationOptions: { pictures: true }, then shapescript/setPictures with {enabled}) only in a trusted workspace, and only while the setting shapescript.hover.pictures is on (it is unless you turn it off); it follows the workspace being trusted and the setting changing. The server runs only functions that only compute: none that print, save, load, write logs, call Python or change a mut parameter. The interpreter is the one go to definition into Python uses: the Python extension's, else shapescript.debug.python, else python3. Without PyTorch there, hovers have no pictures. The first hover of a session can take a few seconds while the worker imports PyTorch; the server starts it when a .ss file opens.

Before and after

A hover on |> or -> into a built-in that only moves values, or on such a built-in, shows what the step changed, and so does the line of types above a picture. The server compares the types before and after (crates/lsp/src/changes.rs, using the pattern's names when the call has one, so einsum's m is the same axis in both operands) and colors the changes, in the types above the image and in the diagrams (the axis's label and dimension line, its part of × H × b, and caption parts):

• orange: what's new, such as softmax's Prob over m, a new axis or dtype;
• red, struck through: what goes away, such as an axis einsum contracts or reduce removes, a domain, or a structure such as −∞ elsewhere;
• blue, purple, green: the same axis before and after, in the same color, when it moved, was grouped (H and Dh before, (H Dh) after), split, resized or taken from another operand (Dh from v).

A step that changes nothing has nothing colored. The extension tells the server it shows HTML (initializationOptions: { hoverHtml: true }, with the client's markdown.supportHtml), so the types line is <code> with colored <span>s, which VS Code's hover sanitizer allows. The server draws the accents as VS Code's default charts.orange, charts.red, charts.blue, charts.purple and charts.green for dark themes, and the extension replaces them, in the image and the spans, with the theme's own charts.* colors, or darker defaults on light themes. Other editors get the drawings as text, with ^ under what changed and ~ under what goes away.

A model's architecture

A model's architecture is drawn in a tab of its own, Architecture: Gpt, beside the editor (VIZ-13). Clicking a model's name, where it's declared (model Gpt:), in a type (m: Gpt) or in Gpt.init, or the name of a function that takes a model first (fn forward(m: Gpt, …), or a call of one), shows that model there. There's one tab, updated in place as the cursor moves to other models, and it never takes the focus from the editor. It opens by itself the first time the cursor lands on a model; once you close it, it stays closed in that workspace until you ask for it: Show architecture above each model line (a CodeLens), Open architecture in the model's hover, Show Model Architecture in the editor's context menu on a model's name, or ShapeScript: Show Model Architecture in the command palette, at the cursor. shapescript.architecture.followCursor: false stops it following the cursor (and opening by itself). In the tab:

• the drawing, full size, in the theme's colors (it follows theme changes), fitted to the tab's width; −, + and Fit zoom, as do Ctrl/Cmd + scroll and the +, - and 0 keys, and scrolling or dragging the background pans;
• Layers and Every operation switch shapescript.architecture.detail;
• the header names the model and the functions it was read from, and those its calls were followed into, each a link that shows its line in the editor;
• clicking a box selects the code it was drawn from: an operation's expression, or a layer's, from its first operation to its last.

The hover, which has no room for the drawing, sums it up in a line built from the layers, Architecture of Gpt: Embedding → Transformer block × 6 → LayerNorm → Linear · tied · Open architecture, then the model's code and docs. The extension tells the server it has the tab (initializationOptions: { architecturePanel: true }); other editors keep the drawing, or its text, in the hover. The tab asks the server for shapescript/architecture with the document and the cursor's position (and the level), and gets the SVG, with each box in a <g data-node="id">, the places of the functions and of each box, and the summary; probe: true asks only whether there's a model there, which sets the context key shapescript.cursorOnModel for the context menu. The page (media/architecture.js, media/architecture.css) is loaded under a content security policy that allows only the extension's own script (with a nonce) and style sheet, from media/, and data: images: no remote content, no inline scripts. The SVG is the server's, with the accents recolored to the theme's charts.* colors (as in hovers) and currentColor left for the page's text color, --vscode-editor-foreground.

The server (crates/lsp/src/architecture.rs) finds the functions that take the model first, in the file and the files it imports, traces forward and loss with the checker's type of each expression, finds the layers in them (crates/lsp/src/architecture/layers.rs), and draws them top to bottom (crates/viz/src/arch/), as papers do: a box per layer, such as Linear D → C, Multi-head attention · H heads · causal or MLP D → Dff → D · GELU, with its result's shape and the params it holds inside it, each as a small rectangle in its proportions; residual connections down the left into a ⊕; dropout as a small tag; and a dashed frame for each function a call was followed into that isn't one layer, or a loop's body, with × L when it repeats, called Transformer block when it holds attention, an MLP and a residual connection.

The setting shapescript.architecture.detail draws every operation instead (operations, rather than layers): each matmul, add and rearrange, with the params beside the operations that read them. The tab's buttons change the setting (in the workspace if it's set there, else for the user), and the extension tells the server (shapescript/setArchitectureDetail). The hover's Open architecture link runs shapescript.architecture.show with where the hover is; the middleware lets that hover run that command alone (isTrusted: { enabledCommands }). Other editors get the drawing in the hover, as a list of the layers or operations as text, with a header that names the functions read, with file:// links to their lines.

Hints

The server sends inlay hints without tooltips, and gives a hint's type and diagram when VS Code resolves it (inlayHint/resolve), which it does when the hint is hovered. The extension's resolveInlayHint middleware draws the diagram in the hover's text color, as for hovers. Clients that can't resolve tooltips get the type as text.

Packaging

A release packages one VSIX per platform, each with its own shapescript-lsp in server/, and publishes them to the Marketplace (not Open VSX): ci/release.sh tags it, and the GitHub mirror's Release workflow does the rest, as docs/releasing.md describes. To package one by hand:

cargo build --release -p shapescript-lsp
mkdir -p server && cp ../../target/release/shapescript-lsp server/
npx vsce package --target darwin-arm64     # or linux-x64, win32-x64, ...

The Marketplace icon, images/icon.png, is drawn from images/icon.svg; after changing the SVG, run rsvg-convert -w 256 images/icon.svg -o images/icon.png. The file icons for .ss files are images/file-light.svg and images/file-dark.svg.