Make the quadrants API type-checkable: annotation spellings, kernel decorator overloads, element-typed containers

[duburcqa]

Concisely describe the proposed feature

Quadrants cannot currently be type checked by a downstream project, for several independent reasons. This issue covers all of them; the annotation spellings below are the largest and the one that has to be fixed first, and a second set of blockers in kernel bodies, decorators and call sites is described further down.

Kernel parameter annotations that are built by calling a factory - qd.types.vector(3), qd.types.matrix(3, 3) - are not valid Python type expressions. A call is not part of the type-expression grammar, so it is rejected regardless of what the callee returns, and the rejection is not tool-specific: mypy 1.15 reports it as [valid-type], pyright 1.1.407 as reportInvalidTypeForm, ty 0.0.65 as invalid-type-form. Any downstream project that wants to run a type checker over code containing quadrants kernels therefore cannot, and this blocks Genesis-Embodied-AI/genesis-world#1271.

Quadrants already solves this for two of its three annotation families. Measuring on Genesis at 54b9ca2e5, ty check genesis reports 5376 diagnostics, of which 1716 are invalid-type-form across 54 files. Grouped by the exact annotation expression:

annotation form	sites	type-expression-valid spelling available today
qd.types.ndarray()	813	yes - qd.types.NDArray
qd.template()	588	yes - qd.Template
qd.types.ndarray(ndim=N) and other keyword variants	53	yes - qd.types.NDArray[None, N]
qd.types.vector(3), qd.types.vector(4)	257	no

So 1454 of the 1716 are already fixable downstream today, and the remaining 257 are blocked here.

The reason the first three work is that qd.types.ndarray and qd.template are aliases for classes (NdarrayType, Template), so the bare alias is a plain class reference - a valid type expression - and NdarrayType additionally defines __class_getitem__, so the subscripted form carries the parameters through. By contrast qd.types.vector and qd.types.matrix are factory functions returning VectorType / MatrixType instances; neither class is exported from quadrants.types nor subscriptable, so for those two there is no valid spelling at all.

Verified against quadrants 1.1.4 (commit 47e9dfb5) and ty 0.0.65 - the following file both runs correctly and reports All checks passed!:

import numpy as np
import quadrants as qd

qd.init(arch=qd.cpu)


@qd.kernel
def bare(x: qd.types.NDArray, s: qd.i32):
    x[0] = s


@qd.kernel
def subscripted(x: qd.types.NDArray[qd.i32, 1], s: qd.i32):
    x[0] = s


@qd.kernel
def templated(f: qd.Template, s: qd.i32):
    f[0] = s


arr = np.zeros(4, dtype=np.int32)
bare(arr, 7)
subscripted(arr, 8)
templated(qd.field(qd.i32, shape=4), 9)

Rewriting any of those annotations into its call form (qd.types.ndarray(), qd.types.ndarray(ndim=1), qd.template()) leaves the runtime behaviour identical and produces error[invalid-type-form] Function calls are not allowed in parameter annotations. There is no equivalent rewrite for qd.types.vector(3).

Describe the solution you'd like (if any)

Export Vector and Matrix from quadrants.types as class aliases carrying __class_getitem__, mirroring exactly what NdarrayType already does, so that the dimensions and dtype travel through a subscript instead of a call:

qd.types.Vector[3]               # equivalent to qd.types.vector(3)
qd.types.Vector[3, qd.f32]       # equivalent to qd.types.vector(3, qd.f32)
qd.types.Matrix[3, 3, qd.f32]    # equivalent to qd.types.matrix(3, 3, qd.f32)

Concretely: export the classes currently named VectorType and MatrixType under quadrants.types as Vector and Matrix, and give each a __class_getitem__ forwarding its arguments positionally to the constructor - the same few lines already present as NdarrayType.__class_getitem__ in python/quadrants/types/ndarray_type.py. Per the naming point below these should be renamed, not aliased, so the class name and the exported name coincide. The existing vector() and matrix() factories can keep working unchanged, which makes the subscript form purely additive on its own.

While settling on a spelling, the two families that already work are themselves inconsistent with each other, and that should be fixed rather than propagated to Vector / Matrix. Three separate things differ:

• Export depth. Template is reachable at both qd.Template and qd.types.Template (they are the same object), whereas NDArray exists only under qd.types - qd.NDArray raises AttributeError. So the two type-expression-valid spellings a user has to write today are qd.Template and qd.types.NDArray, at different depths, for no reason a caller can infer. The top-level copies are also the minority case: vector, matrix, struct, NDArray and BufferViewType are already reachable only under qd.types, and the sole source of the leak is from quadrants.types.annotations import * at python/quadrants/__init__.py:52, whose __all__ is ["template", "sparse_matrix_builder", "Template"].
• A name that means two different things at the two depths. qd.template is qd.types.template, but qd.ndarray is not qd.types.ndarray: at top level it is the allocation function (qd.ndarray(qd.i32, shape=4)), and only under qd.types is it the annotation class. Worse, that trap is invisible to the checker - reveal_type(qd.ndarray) is Unknown because the top-level re-export is not statically resolvable, so def k(x: qd.ndarray) passes ty check with All checks passed! and fails only at import time with QuadrantsSyntaxError: Invalid type annotation (argument 0) of Quadrants kernel: <function ndarray ...>. A user who has learned that qd.Template works at top level and reaches for qd.NDArray or qd.ndarray gets either an AttributeError or a spelling the type checker actively endorses and the runtime rejects.
• Class naming. Four styles coexist. Template is plain CapWords. NdarrayType carries a Type suffix and a different casing of "ndarray" than the NDArray alias in front of it. BufferViewType is exported with its suffix intact, so the use site reads qd.types.BufferViewType. And sparse_matrix_builder is a lowercase class. The suffix is pure stutter under this namespace - qd.types.NdarrayType says "type" twice - so it should come off the classes themselves rather than being papered over by an alias, which also makes the class name and the exported name coincide. The same applies to the classes behind the request above (VectorType, MatrixType) and to StructType and CompoundType, so the new names should land suffix-free from the start rather than being renamed later.

The concrete request is therefore that qd.types become the single home for annotations, with exactly one statically-resolvable spelling per family and no Type suffix: qd.types.NDArray, qd.types.Vector, qd.types.Matrix, qd.types.Struct, qd.types.Template, qd.types.BufferView, qd.types.SparseMatrixBuilder. That means dropping the lowercase duplicates (qd.types.ndarray, qd.types.template) and deleting the from quadrants.types.annotations import * line so that nothing annotation-related remains at top level - which also dissolves the qd.ndarray collision structurally, since the top level would then hold only the allocation function and qd.types only the annotation. The from quadrants.types.primitive_types import * shortcut on the next line should stay as it is: qd.f32 and friends are dtypes rather than annotations, and are genuinely wanted at top level.

This is a breaking change and would need the MAJOR label, but the cost to downstream is close to zero if it ships alongside the Vector / Matrix addition, because every affected call site has to be rewritten anyway to become type-expression-valid - that is one migration instead of two. For the largest consumer it is free: Genesis reaches all of these exclusively through qd.types.* and qd.template(), and imports none of NdarrayType, VectorType, MatrixType, StructType, BufferViewType or CompoundType by name.

One dead end worth recording, so it is not mistaken for a solution: qd.math.vec3 looks like it already covers this case, and a checker does accept def k(v: qd.math.vec3) without complaint. That pass is vacuous. vec3 = vector(3, cfg().default_fp) is a module-level variable whose type a checker cannot infer (the cfg() chain is unresolvable), so it resolves to Unknown and the annotation is simply not checked rather than being valid - reveal_type(vec3) confirms Unknown. Pre-built instances hide the problem instead of fixing it; the spelling has to be a class reference or a subscript of one.

A second blocker: valid annotations are necessary but not sufficient

Fixing the annotations only clears a third of the picture, so it is worth scoping this issue as full typing compliance rather than just parameter declarations. Of the 5376 diagnostics on Genesis, 1716 are the annotations above and a further 1549 sit in the same 58 kernel-bearing files and come from kernel bodies, decorators and call sites. Not all of those are quadrants' responsibility - a large share is Genesis's own (164 from attributes initialised to None and dereferenced without narrowing, 251 from attributes never assigned in their class, and 88 from a Genesis-side name collision where genesis.utils.geom.qd_vec3(val) is star-imported into gs and then shadowed at gs.init() time by qd_vec3 = qd.types.vector(3, qd_float)). The following, though, are quadrants-side and account for 334 of them.

@qd.kernel(fastcache=False) matches no overload - 149 sites. The two @overload declarations for kernel in python/quadrants/lang/kernel_impl.py accept pure, graph and checkpoints, but the implementation below them also accepts fastcache, so the overloads have drifted from the implementation and every parameterised use of the decorator is an error.

qd.kernel erases the signature of the function it decorates, so no kernel call site is checked at all. Both overloads return Any / Callable[[Any], Any] rather than being generic in the decorated callable. qd.func already does this correctly one file over, with F = TypeVar("F", bound=Callable), which is the same kind of internal inconsistency as the naming one above - the fix for kernel is to mirror func.

qd.simt.block.SharedArray defines neither __getitem__ nor __setitem__ - 72 sites. Indexing a shared array (pivot_row[j]) or assigning through it is therefore an error wherever it is used.

There is no way to spell "field or tensor whose elements are T" - 113 sites. Genesis declares struct-element buffers (ContactFace, ContactNormal, MDVertex, EPAPolytopeFace, Witness, ContactHalfspace) and indexes them as gjk_state.contact_faces[i_b, i_n].normal1, but the only available annotation names the element type, which is not subscriptable. This is the same underlying gap as the vector/matrix request: the annotation describes the metadata object rather than the value the kernel body actually receives.

All three of the first points reproduce in this standalone file on quadrants 1.1.4 (commit 47e9dfb5) with ty 0.0.65:

from typing import reveal_type

import quadrants as qd

qd.init(arch=qd.cpu)


@qd.kernel
def plain(s: qd.i32) -> qd.i32:
    return s + 1


@qd.kernel(fastcache=False)
def with_fastcache(s: qd.i32) -> qd.i32:
    return s + 1


@qd.func
def helper(s: qd.i32) -> qd.i32:
    return s + 1


reveal_type(plain)   # Any            <- signature erased
reveal_type(helper)  # def helper(s: Unknown) -> Unknown
plain("this is not an i32", 1, 2, 3)  # no diagnostic; the runtime does reject it

ty check reports exactly one error on it - error[no-matching-overload] No overload of function 'kernel' matches arguments on the fastcache line - while with_fastcache(4) runs fine, confirming the overloads and not the implementation are wrong. Swapping fastcache=False for graph=False removes the error, isolating it to that one missing parameter. The last line shows the consequence of the erased signature: four arguments of the wrong types to a one-argument kernel pass the checker silently, and only the runtime catches it.

The remaining 712 diagnostics in those files I have not attributed yet, and I have not checked whether they collapse to a similarly small number of root causes the way the annotations did.

Additional comments

Separately, the distributed package carries no py.typed marker and no .pyi stubs, so under PEP 561 a conforming checker is not supposed to use its inline types for a normal (non-editable) install at all. That is a different question from this one - whether quadrants wants to advertise types downstream - so I have not folded it in here, but happy to open it as its own issue if that is useful.

[D0liphin]

Hey, we've been working on a library that provides some typing for Quadrants. We arrived at some of the same conclusions as you :)

The library is available here. It was quite a bit of work to get even what we have here working, despite it not looking like much (!!). Note that it is far from complete.

Differences in Goals

It's important to highlight that this library has somewhat different goals:

1. We only care that the annotations provide a correct set of typing restrictions and are easy for the programmer to reason about. We don't necessarily care that they are annotations of existing Quadrants constructs. This doesn't seem too far from what you are proposing.
2. We want the types to be fully checkable and display nicely in your linter, including tensor rank and matrix/vector dimensions.
3. We disallow int/float promotion and require explicit casts (more on this below).

Overview

We take a very similar approach to the one you've already shown here. We call it qdt (similar to numpy.typing as npt)

Description	qdt version	qd version	Example
Concrete tensor type	qdt.NDArray[qd.<dtype>, Dim<N>]	qd.types.NDArray[qd.<dtype>, <N>]	qdt.NDArray[qd.i32, Dim3]
Templated tensor type	qdt.NDArray[T, DimAny]	qd.types.ndarray(dtype=None, ndim=None)	def f[T: (qd.i32, qd.f32)](arr: qdt.NDArray[T, DimAny])
Vector type	qdt.Vec[qd.<dtype>, Dim<N>]	qd.types.vector(<N>, qd.<dtype>)	qdt.Vec[qd.f32, Dim3]
Matrix type	qdt.Mat[qd.<dtype>, Dim<R>, Dim<C>]	qd.types.matrix(<R>, <C>, qd.<dtype>)	qdt.Mat[qd.f32, Dim3, Dim4]
Compile-time value	qdt.Template[T]	qd.template()	qdt.Template[bool]

In addition to the above, we have only typed some basic functions, including qd.kernel and qd.func, as well as qd.simt (including SharedArray) and qd.math.

Design Choices

qdt.DimN instead of integers

Instead of qd.types.NDArray[qd.f32, 4], we write qdt.NDArray[qd.f32, qdt.Dim4]. The benefit is that this lets you type things based on the dimension. We find that we never need anything above four dimensions, but we provide types up to Dim8.

An example of where this is useful is when we type __matmul__ (_Num is a TypeVar). Here we're only handling square matrices, but we should really support the general case. This should be an easy fix later :)

@overload
def __matmul__[D_: DimAny](
    self: Mat[_Num, D_, D_], other: Mat[_Num, D_, D_], /
) -> Mat[_Num, D_, D_]: ...

@overload
def __matmul__[D_: DimAny](
    self: Mat[_Num, D_, D_], other: Vec[_Num, D_], /
) -> Vec[_Num, D_]: ...

Alternatives

• Use aliases to literals, e.g. Dim1 = Literal[1]. I think this is mostly better, except that your linter might display qdt.NDArray[qd.f32, Literal[4]], which is a little ugly :/
• Do not include dimension information. I'm not a fan of this. I think it's useful to know that you can index a qdt.NDArray[T, qdt.Dim3] with a qdt.Vec[int, qdt.Dim3] :)

Integer Promotions

Since we were adding these types for our own use, we decided to disallow integer/float promotion. Initially, this meant we also didn't allow float or int in any Quadrants operations. This was useful because it gave us a clean set of fixed-width integer types. However, it also made code quite cumbersome:

qd.ndrange((qd.i32(0), x), (qd.i32(0), y), (qd.i32(0), qd.i32(3)))

Instead, we decided to type int as i32 all the time, and potentially release different versions for different widths. The current and only version would be called quadrants_typing_i32_f32 when we adopt this approach.

This creates a somewhat awkward situation where it's hard to constrain types properly. We can't define a qdt.Integer because we can't make int inherit from it. So how would you constrain a generic T to integers only?

You couldn't do T: (qdt.Integer, int) or T: qdt.Integer | int. The former will assign T as qdt.Integer, which is not operable with int, and the latter allows T to be u1 | int, which wouldn't make sense when used in qdt.NDArray[qd.u1 | int, qdt.Dim3], for example.

The only correct way is to explicitly list all possible types for T. So if you really want all integers, you'd have to write:

@qd.kernel
def mykernel[T: (u8, u16, u32, u64, i8, i16, i32, i64)](
    x: qdt.NDArray[T, qdt.DimAny]
) -> None: ...

:/ !!

This wouldn't be especially difficult to change, although I would personally advocate that disallowing promotions is much nicer than allowing them. At least within our team, it makes things considerably simpler :)

Main Limitations

1. There is no way to enforce that someone types a kernel correctly, i.e. that they use only parameter types that are valid for kernels. This is not a high priority, since it would be caught on call or even during decoration, but it is still not ideal. This is hard to do using pure types, but it could be checked at runtime in the decorator, which would be almost as good.
2. There is no good way to construct Quadrants primitive types in Python-land, so you can't accurately construct the types needed to call a kernel. We currently handle this internally using no-op functions such as as_i32(n: int), but that isn't ideal. Could we instead define our own qd.i32 type in Python-land as a plain subclass of np.int32, and similarly for the other types? 🤔
3. It is still very incomplete. We've implemented the main data types described above, along with simt and math, but that's it for now :)

[hughperkins]

@D0liphin Very nice!

Question, for "There is no good way to construct Quadrants primitive types in Python-land", to what extent does the work in the following PR address this issue for you: #414

[D0liphin]

@hughperkins #414 completely solves this aaict. Thank you!

[hughperkins]

@hughperkins #414 completely solves this aaict. Thank you!

Great! 🎉

Dont suppose, do you have a moment to validate/review that PR? Like, no need to look at the code (though you can), but at least confirm in detail that it does exactly what you need, and raise any concerns/questions you have about any deltas between what it does and what you need. Also, check that, if we merged it now, as-is, would it break any of your existing work?

[hughperkins]

By the way @D0liphin what is your preferred future for your library. Do you prefer that:

• it stays a standalone library?
• its features are merged into Quadrants?

(I haven't reviewed your library, at this point, but realized i should probably check this top-level question first)

[D0liphin]

@hughperkins thanks so much for the reply and I apologise for my poor reply speed.
I would prefer if we could merge this kind of thing into quadrants.types, however, there are some things missing from it at the moment. We use it internally for all our typing needs and it does well!
I am happy and eager to do the work to get this kind of thing through, but I would first like to iron out some of my more opinionated decisions (e.g. float being equivalent to f32), so it probably remains experimental (and separate) for another little while, as I iron those out.
if you could take a look at the documented design decisions above, and then also the examples folder of the repository, and let me know what you think, that would be very helpful!

[hughperkins]

Yes, anything opinionated / experiimental, lets try to keep separate for now please.