Python interop

import python torch binds the Python module torch. Members of a Python module (torch.load), calls to them, and anything computed from a Python value are Python values. A Python value has no type: the checker knows nothing about it, not even whether it is a tensor, and arithmetic or a comparison with one gives another Python value. A Python value can be:

• passed back to Python, or combined with other values;
• passed for a shapeless parameter or a host scalar parameter, which takes it unchecked, or for a model, Mode, optimizer state or Schedule parameter, which checks it at run time; it can be a learning rate or a hyperparameter of an optimizer, the model or state of update, the gradient of clip_norm, and the path or a value of save or the path of load, which save checks when it runs, but not the model of sgd or adamw, whose state’s type comes from the model’s (E1501);
• an if condition, an end of a range, or what a for loops over, whose items are Python values too (Control flow);
• given a type with as: a tensor type such as as [D 3*D], or as int, as float or as bool (Host scalars).

import python torch

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

Using a Python value where a shape is needed, such as returning it from a function with a tensor return type, passing it as a tensor argument, giving it to a pattern operation or to an operation along a named axis, or converting it with to and a domain, is E0413.

In the other direction, every ShapeScript function checks its arguments against its signature when it is called, whether from Python or from ShapeScript (Runtime checks).

In the editor, go to definition on the module of an import python (on any part of torch.nn, up to that part), on a member of one (torch.load, augment.random_crop), or on an extern fn’s path string opens the Python source. The language server asks a Python interpreter: the python initialization option or the last shapescript/setPython notification’s {python} (the VS Code extension sends the Python extension’s selected interpreter, else the shapescript.debug.python setting), else python3 (python on Windows). It runs a small script that finds the module with importlib.util.find_spec, looking in the file’s folder and the workspace’s folders first, then finds the name in the module’s source with Python’s ast, without running the module; a name the module only imports, as torch does load, is found by importing the module and asking inspect, and one defined in C by the package’s .pyi stubs. Answers are kept for each interpreter and name until the server restarts.

The other way, the language server answers the request shapescript/declarationOf with {path, name}, the location of the top-level declaration name in the .ss file path (the first part of a dotted name such as Mlp.init; an empty name is the file itself), which the VS Code extension uses for go to definition in Python files (Python modules for .ss files).

Python modules for .ss files

Python can import a .ss file in two ways: with the import hook, after shapescript.install(), or through the Python module generated beside it, model.py for model.ss. The module is what Python tools read: Pylance, pyright and mypy know nothing of .ss files, so without it from model import f is an unresolved import to them, even when it runs. It has:

• a first line that says it’s generated: # Generated by ShapeScript from model.ss. Edit model.ss; this file is rewritten when it changes.;
• a typed declaration of everything the compiled file defines, each with a docstring holding its ShapeScript declaration and its ## comment, which editors show on hover;
• a last line, __import__("shapescript.pymodules").pymodules.load(globals(), __file__), which compiles model.ss, with the import hook’s cache in __pycache__ and the check modes given to install() if it was called, and puts the compiled dims, models and functions in place of the declarations. So import model works without shapescript.install(). As with the hook, the code has the .ss file’s lines, so tracebacks and debuggers show ShapeScript, and a file with errors raises ShapeScriptError from its first error’s line (see Debugging).

The module imports torch and typing as _ss_torch and _ss_typing, so the table’s torch.Tensor is written _ss_torch.Tensor: that way the only names the module offers, as after from model import in an editor, are the file’s own (editors show the types themselves, Tensor). Tensor types become torch.Tensor; host scalars int, float or bool; a model, its class; Mode, Grad[M], optimizer state, Schedule, Dataset[R] and batches, the classes of shapescript.rt; and a LogSink, typing.Any. Parameters keep their order and defaults, as in the compiled function. Shape aliases are only for types, so they aren’t declared.

A file that imports others is checked with them, found next to it and then on sys.path (from Python) or $PYTHONPATH (the command line). Its module doesn’t import theirs, which may not have been generated: another file’s model is torch.nn.Module, its data record’s batches "_ss_rt.Batch" and its logs record "_ss_rt.Logs". When Python imports an imported file through its own generated module first, the importing file uses that module rather than compiling the file again, so models and records are the same classes. Overloaded functions are one declaration of the name, forward: _ss_overloads_forward, typed by a typing.Protocol whose __call__ has each as a typing.overload with its ShapeScript signature, so that editors find one place for the name (the variable’s docstring lists every signature); overloads that Python’s types can’t tell apart, such as two that take models of other files, both torch.nn.Module, are declared once, with every ShapeScript signature in the docstring, since type checkers reject an overload that can never be picked. A generic function is declared once under its name, which is the function that calls the instance for its argument’s class: a type variable’s model is torch.nn.Module, and its record’s batches "_ss_rt.Batch".

With the hook installed, the hook imports model.ss itself (it comes before Python’s own finder), so the module’s declarations are only read by tools, and the file is compiled once either way.

The language server writes the module when a .ss file is opened or saved and has no errors; the VS Code setting shapescript.pythonModules turns this off. shapescript pymodules FILE... writes them from the command line, and shapescript pymodules --check FILE... only fails if one is missing or out of date, for CI. The import hook rewrites a generated module that is older than its file. None of them ever replaces a .py file that doesn’t start with the generated line: that is someone’s own module, and the language server and pymodules report W2501 instead.

Going between Python and ShapeScript. In VS Code, Ctrl-click (go to definition) on a name a Python file imports from a .ss module goes to its declaration in the .ss file, never to the generated module, and so does Ctrl-click on a declaration in the generated module itself. Pylance, as any Python language server, answers with the generated model.py, and VS Code shows every extension’s answers, with no way for one to drop another’s; so the extension adds no answer of its own when another one answers, and when an editor then lands where that answer said, it shows the .ss file at the same declaration in its place, in the same tab. Going back from there skips the generated module. Without a Python language server, the extension’s own answer is the .ss file. A generated module opened on purpose, from the explorer or its tab, stays open: only an editor that lands where a go to definition just said is replaced. The module is looked for as Python would find it, next to the Python file, then in the workspace’s folders; a name of a module that is only a .ss file (which the import hook imports) goes to it too. The other way, the references of a .ss file’s top-level declaration (Find All References, and what Ctrl-click on the declaration shows) include the Python code that uses it: the uses Pylance finds of the generated module’s declaration, leaving out the generated modules, or, without a Python language server, the uses in the workspace’s Python files that import the module, read from their imports (from model import f as g, import model as m) and text, not their scopes. Peek Definition still shows Pylance’s answer, the generated module, as a peek doesn’t open an editor.

The modules are derived files. A project can ignore them (*.py beside .ss files in .gitignore), and they’re written again whenever the editor opens a file; but a type check in CI then needs shapescript pymodules first. Or a project can commit them, so that import model type-checks for anyone, with shapescript pymodules --check in CI to keep them up to date. Either way works; ignoring them keeps diffs to the .ss files.