shapels

Shape inference for PyTorch tensors directly inside VS Code. The extension bundles the shapels language server binary for Linux, macOS and Windows so users can get diagnostics without extra setup.

Features

• Reports tensor shape inference diagnostics as you type.
• Runs inference for tensor shapes, allowing for hovering over variables to inspect their shape.

Requirements

No external dependencies. The platform-specific server binary is shipped inside the extension (shapels-bin/<platform>/shapels).

Extension Settings

• shapels.serverPath: Override the bundled server binary path.

Release Notes

0.4.1

New features and operations:

• Inference now propagates through self.attribute, if a tensor, parameters/buffers or a torch.nn.Module.
• Support for common native torch.nn.Module: Linear, BatchNorm*, Conv*, Pool*, ConvTranspose*, Sequential and Flatten.
• Support for losses, functional and as torch.nn.Module.
• Calling annotated functions supports renaming and ellipsis are properly applied ()
• Infer annotated Callable.
• torch.new_* creation ops, like new_zeros.
• Infer torch.take, torch.where and torch.logaddexp.
• UnaryOp, FloorDiv, Pow and Mod operator inference.

Enhacenments and fixes:

• Session-level caching for modules, improving performance quite drastically.
• Run inference through with contexts.
• Negative args for tranpose/permute are properly handled.
• masked_fill and quantile are more generally inferred with a custom implementation.
• Matmul in batch matrix mutiplication cases is now properly inferred.
• torch.arange with a single argument now behaves correctly.
• Boolean masks as indices now resolve to a range, since they cannot be absolutely known at runtime.

0.4.1

Minor bug fixes:

• Improved diagnostics on some calls using non-existing dimensions of a tensor.
• Removed spurious diagnostics on caller site for not understood dimensions on the callee.

0.4.0

• Inference now propagates through calls to instances of torch.nn.Module to the respective forward method.
• Inference also propagates across non-forward methods for caller–callee analysis.
• Inference runs through control flow (if, for, while) rather than skipping branches.
• Functional convolution ops are supported.
• New broadcastable operation groups: functional elementwise ops (torch.add/torch.div etc.) and broadcastable comparisons (with boolean dtype output).
• Bitwise operator support added for tensor operations.
• New repeat/expand family support: torch.expand, torch.repeat, and torch.repeat_interleave.
• View/reshape ops expanded: torch.flatten, torch.ravel, plus spatial transforms flip and rot90.
• Creation op coverage widened with *_like ops and range/lin/logspace (plus randperm) creation.
• Partial support added for masked_fill.
• New CLI to analyze a single file and print results to stdout.

0.3.0

• Interactive tutorial at https://carrascomj.github.io/shapelsbook/
• Support for creation ops like tensor.zeros, tensor.ones and shape unrolling! This allows for not using jaxtyping at all if desired.
• Support for indexing.
• Support for comparison operators.
• Support for .to, and other operations like bool, int, float, etc.
• Support for randn, rand, detach and cuda.