pyrefly-warp-stubs

NVIDIA Warp type stubs with array and tile shape tracking for Pyrefly, in the spirit of Pyrefly's pyrefly-numpy-stubs and pyrefly-torch-stubs.

@wp.kernel
def mlp(x: wp.array2d[float], w: wp.array2d[float], out: wp.array2d[float]):
    inputs = wp.tile_load(x, shape=(16, 32))      # Tile[float, [16, 32]]
    weights = wp.tile_load(w, shape=(64, 16))     # Tile[float, [64, 16]]
    hidden = wp.tile_matmul(weights, inputs)      # Tile[float, [64, 32]]
    wp.tile_matmul(inputs, weights)               # error: expected Int[32], got Int[64]
    bias = wp.tile_load(w, shape=(64, 1))
    hidden + bias                                 # error: tiles never broadcast implicitly
    wp.tile_sum(hidden, axis=2)                   # error: reduction axis out of bounds

Tile shapes in Warp are compile-time constants, so a shape bug normally only surfaces when the kernel is built at launch. With these stubs Pyrefly reports it in the editor, along with Warp's scalar and vector type rules (wp.max(f64, 0.0), vec3 * vec3, mat44 * vec3 are errors, as in Warp).

Install

Add it as a development dependency of a project that uses Warp:

uv add --dev pyrefly-warp-stubs        # or: pip install pyrefly-warp-stubs
# latest main: uv add --dev "pyrefly-warp-stubs @ git+https://github.com/julien-blanchon/pyrefly-warp-stubs"

This installs warp-stubs (a PEP 561 partial stub package) and pyrefly-shape-extensions. Pyrefly picks the stubs up automatically. Requires Python ≥ 3.12 and Pyrefly ≥ 1.4, in default or strict mode: add a pyrefly.toml (or [tool.pyrefly] in pyproject.toml), since the basic mode used without one hides shape errors.

The stubs are generated from warp-lang 1.18. Newer Warp releases install fine, but the stubs replace Warp's own __init__.pyi, so functions added after 1.18 are reported as missing until the stubs are regenerated (see Development).

Using it in your project

The stubs only exist for the type checker: Python never imports .pyi files, so nothing changes at runtime, and people who install your package never download these stubs or need Pyrefly. That holds as long as shape annotations stay type-checking-only:

from __future__ import annotations  # required: see below

from typing import TYPE_CHECKING

import warp as wp

if TYPE_CHECKING:  # never imported at runtime
    from shape_extensions import IntVar


def normalize[N: IntVar](points: wp.array[wp.vec3, [N]]) -> wp.array[wp.vec3, [N]]:
    ...

normalize(wp.zeros(10, dtype=wp.vec3))           # ok: array[vec3f, [10]]
normalize(wp.zeros((10, 3), dtype=float))         # error: array[float, [10, 3]]

• Keep from __future__ import annotations in every file with shaped annotations. Without it Python evaluates wp.array[wp.vec3, [N]] at import time, and Warp rejects anything but an integer ndim there.
• Import shape_extensions only under TYPE_CHECKING. Using assert_shape, shape_vars or a top-level from shape_extensions import ... in shipped code makes pyrefly-shape-extensions a runtime dependency.
• No shapes in @wp.kernel/@wp.func signatures or @wp.struct fields. Warp evaluates those annotations itself, even with the future import; keep wp.array2d[float] there. Shapes are still inferred inside the kernel body.
• Shape assertions inside kernels go under if TYPE_CHECKING:. Warp constant-folds the condition and never compiles the block, while Pyrefly checks it. A bare reveal_type/assert_type in a kernel breaks the build. expect_shape (from the stub-only warp.shape_checks) checks against your tile-size constants instead of literals:

  t = wp.tile_load(a, shape=(TILE_M, TILE_N))
  if TYPE_CHECKING:
      expect_shape(t, (TILE_M, TILE_N))
      assert_type(t, Tile[float, [8, 4]])
  

• Quote shaped types in executed assert_type/cast: assert_type(x, "wp.array[float, [10, 3]]"). Their type argument is evaluated when the line runs, even with the future import. Shaped type aliases must use PEP 695 (type Points = wp.array[wp.vec3, [int]]).
• Tile, Array, Vector, warp.shape_checks, ... exist only in the stubs: import and use them under if TYPE_CHECKING: or in annotations.
• Public API. Other type checkers (Pyright, mypy, ty) cannot read wp.array[wp.vec3, [N]], so a downstream user without these stubs sees an untyped parameter. If that matters, use the portable Shaped[wp.array[wp.vec3], "[N]"] inside a @shape_vars("N") function, which other checkers read as plain wp.array[wp.vec3]. That decorator runs at import time, so it does make the small, pure-Python pyrefly-shape-extensions a runtime dependency of your package.
• Contributors using other checkers. Pyright or ty in an environment with your dev dependencies will also pick up warp-stubs and report its shape syntax as errors. This only affects people developing your project.

skills/add-shape-types-to-warp-project/ contains an agent skill that applies all of this step by step, and verify_port.py, a lint for these runtime rules (Pyrefly cannot see them).

Shapes on kernel inputs: typed launchers

Kernel signatures must stay plain Warp, and wp.launch(inputs=[...]) is not checked against them. Put the shapes on the host function that launches the kernel; Pyrefly then checks every call, and the dimensions flow to the result:

def matmul[M: IntVar, K: IntVar, N: IntVar](
    a: wp.array[float, [M, K]], b: wp.array[float, [K, N]]
) -> wp.array[float, [M, N]]:
    c = wp.zeros((a.shape[0], b.shape[1]), dtype=float)   # array[float, [M, N]]
    wp.launch_tiled(gemm_kernel, dim=..., inputs=[a, b, c], block_dim=64)
    return c

matmul(wp.empty((64, 32)), wp.empty((32, 48)))   # array[float, [64, 48]]
matmul(wp.empty((32, 48)), wp.empty((64, 32)))   # error: expected [48, @_]

Inside the kernel, a, b and c are only array[float, [int, int]]: a kernel cannot state that two parameters share a shape (a generic kernel signature fails when Warp evaluates it), and inputs=[...] is not checked against the kernel. Allocating outputs from the inputs in the launcher (wp.zeros((a.shape[0], b.shape[1])), wp.empty_like(x)) enforces the relationship by construction and leaves that one line as the only unchecked link.

Named dimensions (B C H W)

Named dimensions are type parameters. Three spellings check the same way:

Arithmetic on them is tracked: x.reshape((x.shape[0], x.shape[1], -1)) on a [B, C, H, W] array is [B, C, H * W] (see examples/image_batch.py).

Libraries: a typing extra

If your library uses Warp and wants its users to get the shape types too, offer the stubs as an extra, not only a dev dependency: dependency groups (PEP 735) are never published, extras are.

[project.optional-dependencies]
typing = ["pyrefly-warp-stubs>=0.1"]       # users opt in: pip install mylib[typing]

[dependency-groups]
dev = ["mylib[typing]", "pyrefly>=1.4", "pytest"]   # contributors get it by default

A plain pip install mylib still installs only Warp. Keep the published requirement a plain name (pyrefly-warp-stubs>=0.1): build backends refuse git URLs in optional-dependencies and PyPI rejects them. To develop against the latest main instead, add a source that is not published:

[tool.uv.sources]
pyrefly-warp-stubs = { git = "https://github.com/julien-blanchon/pyrefly-warp-stubs" }

What is tracked

Literal sizes, module-level constants (TILE = 64) and wp.constant(64) are tracked; sizes only known at runtime become int dimensions, which keep the rank.

Every operator and rule above was checked against what Warp 1.18 actually builds, by compiling kernels on the CPU; the tests keep it that way.

Examples

Each example runs on the CPU (python examples/<name>.py), checks its result against NumPy, and type-checks cleanly; its docstring lists errors Pyrefly reports for typical mistakes (all verified by the test suite or by hand).

How it works

Pyrefly treats any class with a type parameter bound to IntTuple as a shaped array. Warp already declares Array(Generic[DType, NDim]) and Tile(Generic[DType, Shape]). These stubs re-declare them as Array[DType, Shape: IntTuple] and Tile[DType, Shape: IntTuple], keeping the shape second so that wp.array[float] still means "dtype float".

Warp's __init__.pyi is ~11k autogenerated lines, and the classes live in warp/_src, so the stubs are generated rather than hand-written:

overrides/<path>.pyi    ─┐                          ┌─> warp-stubs/<path>.pyi
installed warp package  ─┴─ tools/generate_stubs.py ┘   (_shapes.pyi, shape_checks.pyi are hand-written)

Each file under overrides/ stubs the Warp module at the same path (12 modules: __init__, _src/types, renderers, sparse, several fem modules). tools/generate_stubs.py copies the module and merges in the overrides:

• An override def replaces every upstream overload of that name; the upstream docstring is carried over (also for functions upstream only re-exports).
• An override class merges into the upstream class: it supplies the header and the members it redefines, and keeps the rest.
• rewrite_params widens upstream parameter annotations that are narrower than what the code accepts (e.g. renderer pos: tuple also takes a wp.vec3).
• Warp's @wp.func overloading by redefinition becomes typing.overload groups; otherwise the last definition would shadow the others.
• An override import may not silently rebind a name upstream imports from elsewhere (an early version replaced wp.overload with typing.overload).
• Function bodies are kept: Pyrefly infers attributes and return types from them even in stubs.

Shape computations that need logic (reshape, axis reductions, broadcasting, tile_arange) are Pyrefly shape DSL rules in warp-stubs/_shapes.pyi.

Fixes to upstream typing

Precise element types expose gaps that were hidden while a[i] was untyped, so the overrides also fix these:

• Scalars: float/int/bool and wp.float32/wp.int32/wp.bool were distinct types, so wp.max(a[i], 0.0) and passing a dtype=float array to a wp.float32 parameter were errors. Scalar operators are typed.
• Vectors: upstream has two unrelated vec3f classes (warp/__init__.pyi and warp._src.types), neither a Vector, with constructors that reject valid calls (wp.vec3(np_array)). They are now one type with Warp's operators; wp.types.vector(...) and wp.spatial_vector are vectors too.
• wp.tid() is int | tuple[int, int] | ... upstream, which cannot be unpacked as i, j = wp.tid(). It is now an int that can also be unpacked.
• warp.fem: fields, domains, operators and integrands were not callable (360 errors on Warp's FEM examples), and @integrand was typed as a union.
• Kernel-only fields that Warp declares solely in its code-generation registry (MeshQueryPoint.face, Texture3D.depth, ...) are read by the generator and declared; dtype-driven builtins (texture_sample, volume_sample, cast) return the requested type instead of Any.
• for i in wp.hash_grid_query(...) is iterable; wp.where/wp.select keep the value type; tile_*_solve_inplace accept 2D right-hand sides; tile_fft and the *_indexed functions accept any element type; wp.indexedarray1d[...]–4d[...] work in annotations; fem.Grid2D/Grid3D and the renderers accept the vector values Warp's examples pass them.

On Warp's own 92 example scripts, Pyrefly reports 370 errors with these stubs versus 2145 without. Nine errors are new: one real bug in the examples (range(num_bodies / TILE_SIZE) passes a float), four where example_diffusion_3d.py passes 3D bounds to a 2D helper annotated wp.vec2, and four from a Pyrefly bug (below). See tools/compare_examples.py.

Known limitations

• wp.array[dtype, Literal[ndim]] cannot be expressed: the second parameter is now the shape. Use wp.array2d[dtype] etc.
• Tile shapes cannot be written in kernel signatures or kernel-local annotations: Warp evaluates kernel annotations at runtime and does not support annotated assignments in kernels. Shapes are inferred inside kernel bodies.
• wp.array[float] (no rank) is an array of unknown rank; use wp.array1d[...] to get rank checks.
• wp.launch(kernel, inputs=[...]) arguments are not checked against the kernel signature.
• wp.tile[float, M, N] annotations on @wp.func parameters are valid Warp, but Pyrefly 1.4 cannot express a name that is both the wp.tile(x) builtin and a variadic shape type, so those parameters are Unknown (as upstream).
• Module constants cannot be used as dimensions in annotations (M = 4, Tile[float, [M, 8]]); no spelling works in Pyrefly 1.4. Use literals, IntVar parameters, or expect_shape(t, (M, 8)). Constants are tracked in calls (shape=(M, 8)).
• Typing lets int pass where float is expected, so wp.int32 (which is int) tiles reach float-only tile operations, and wp.max(f32, 1) passes; Warp rejects both. Other integer types (wp.int64, wp.uint32, ...) are caught.
• A list literal passed inline as an argument (f(wp.array([0.5, 0.5]))) is typed from the expected parameter and not checked against it; assigned to a variable first, it is.
• .w on a 3-vector is not caught; the length from wp.types.vector(length=5, ...) is widened to int.
• @wp.func overloading by redefinition in your code: the checker sees the last definition only (the generator fixes this for Warp's own modules).
• Results of tile_map, of field evaluations in warp.fem, and of builtins that upstream declares as Vector[Scalar, Any] (wp.normalize) are not length- or dtype-tracked.
• wp.spatial_vector is a function in upstream's stubs (the kernel constructor), so it cannot be used as a type; use wp.spatial_vectorf.
• Pyrefly 1.4 issues found while building this:
  • A shape list spelled with builtins.int ([builtins.int, builtins.int]) silently becomes Unknown. Bare int works.
  • Calling an unspecialised generic alias (self.x = wp.array2d(data)) reports the alias's type variable as leaking (invalid-type-var), even with a default=Any.

Development

uv venv -p 3.13 && uv pip install pyrefly==1.4.0 pyrefly-shape-extensions==1.4.0 \
    "warp-lang==1.18.*" numpy pytest

.venv/bin/python tools/generate_stubs.py    # after editing overrides/
.venv/bin/python -m pytest tests/           # static + runtime suites
.venv/bin/python tools/compare_examples.py  # errors added/removed on Warp's examples

Each test file pairs static check_* functions (assert_type, with an # E: <message> comment on lines that must error, checked by pyrefly check --expectations) with runtime test_* functions. The runtime functions launch the same operations on Warp's CPU backend: well-typed kernels must match NumPy, and every kernel Pyrefly rejects must also fail to build in Warp. tests/test_static.py runs the Pyrefly half, checks that the generated stubs are up to date, and validates the DSL rules; tests/test_verify_port.py pins the lint.

To target another Warp release: install it, regenerate, and run the tests. Overrides of functions whose upstream signature changed will show up as test failures or new errors in compare_examples.py.

License

The generated stubs contain portions of Warp's source, stubs and docstrings (Apache-2.0, © NVIDIA). warp-stubs/_shapes.pyi adapts rules from pyrefly-numpy-stubs (MIT, © Meta Platforms).