Support optional non-template kernel arguments that specialize away the absent code path

[duburcqa]

Concisely describe the proposed feature

Quadrants should support optional non-template kernel arguments spelled T | None, defaulting to None when unset, and should specialize the compiled kernel on presence versus absence so that the unused code path is never traced. Concretely, this should be valid and supported:

@qd.kernel
def apply_at(out: qd.types.NDArray[qd.f32, 2], point: qd.Tensor | None = None):
    for i in range(out.shape[0]):
        if qd.static(point is not None):
            ...  # use point[i]
        else:
            ...  # use the default anchor

The motivating case is the application point of a force applied to a geometry. The point can be left unspecified (the center of mass is used), given as an absolute position, or given as an offset relative to the center of mass. Each variant wants a different arithmetic path, and the unspecified variant wants no arithmetic at all rather than adding zero or multiplying by one. Today the choice is between duplicating the kernel once per variant, which multiplies the maintenance burden and lets the copies drift, or keeping a single kernel that pays for arithmetic it does not need in the common case.

The real-world instance of this is Genesis-Embodied-AI/genesis-world#3143, which adds an application point to the external force API and hits exactly this trade-off in the kernel behind it.

What works today

The underlying mechanism already exists, but it is not spellable the way the code above reads, and it is not covered by any test.

A qd.template() slot accepts None and specializes on it, since a template argument is part of the specialization key by construction. That is not a usable answer here: the same slot rejects an ndarray outright with Ndarray shouldn't be passed in via 'qd.template()' (python/quadrants/lang/_template_mapper_hotpath.py:268), so an argument that is a tensor when present cannot go through it. The only route that accepts both forms today is qd.Tensor.

qd.Tensor is a value-dispatched argument slot: _extract_arg routes ndarray-shaped values through the ndarray feature path and everything else through the template path (python/quadrants/lang/_template_mapper_hotpath.py:237), so a None argument lands in the specialization key. FunctionDefTransformer._decl_and_create_variable then binds the kernel-scope name to the literal Python None (python/quadrants/lang/ast/ast_transformers/function_def_transformer.py:82), and qd.static folds the branch at AST build time.

So the following runs correctly on CPU at 47e9dfb5:

import numpy as np

import quadrants as qd

qd.init(arch=qd.cpu)


@qd.kernel
def add_bias(out: qd.types.NDArray[qd.f32, 1], bias: qd.Tensor):
    for i in range(out.shape[0]):
        if qd.static(bias):
            out[i] = qd.f32(i) + bias[i]
        else:
            out[i] = qd.f32(i)


out = qd.ndarray(qd.f32, shape=(4,))
bias = qd.ndarray(qd.f32, shape=(4,))
bias.from_numpy(np.full(4, 100.0, dtype=np.float32))

for arg in (None, bias, None, qd.Tensor(bias)):
    add_bias(out, arg)
    print(f"{type(arg).__name__:8s} -> {out.to_numpy()}")

print("specializations:", len(add_bias._primal.mapper.mapping))

NoneType -> [0. 1. 2. 3.]
ScalarNdarray -> [100. 101. 102. 103.]
NoneType -> [0. 1. 2. 3.]
Tensor   -> [100. 101. 102. 103.]
specializations: 2

Exactly two specializations, correct results on every call, and no thrashing when alternating between the two forms. The absent branch is genuinely never traced rather than eliminated later: replacing its body with bias.no_such_attribute[i] still compiles and runs when None is passed.

The workaround does not extend to external arrays

An argument that has to accept a torch tensor when present cannot use any of this. The two annotations that could carry such an argument fail on opposite ends, and there is no third option covering both:

annotation	None	torch.Tensor
qd.types.ndarray(dtype=qd.f32, ndim=1)	QuadrantsRuntimeTypeError: Invalid type for argument bias, got None	works
qd.Tensor	works	ValueError: Input to qd.static must be compile-time constants or global pointers, instead of <class 'torch.Tensor'>

A torch tensor reaches an qd.types.ndarray() slot through the external arrays path, which reads shape = getattr(arg, "shape", None) and raises as soon as that comes back None (python/quadrants/lang/_template_mapper_hotpath.py:335-338). None is therefore indistinguishable from any other invalid argument at that point and never reaches the specialization key.

Conversely torch.Tensor is neither an Ndarray nor an AnyArray, so a qd.Tensor slot falls through to the template path and binds the raw torch object as a template value, which qd.static then refuses (python/quadrants/lang/impl.py:1331). A second problem sits behind that error: the specialization count for None plus two distinct torch tensors plus one quadrants ndarray comes out at four, because each torch tensor object becomes its own key. Even with the static check relaxed, that path would recompile the kernel once per tensor instance.

There is also no conversion route out of this. from_torch copies into an already allocated quadrants ndarray, and dlpack support is export only (to_dlpack), with no import side, so a torch tensor cannot be adopted as an Ndarray and then routed through the qd.Tensor slot.

So the mechanism described above covers exactly one of the two argument families that can carry a tensor, and callers holding external arrays are left with the duplicate the kernel or pay for the no-op arithmetic choice.

Why this is not satisfactory

Everything above is reached by working around the type system rather than through it.

gap	current behaviour	reference
qd.static(x is not None)	rejected with Operator "is not" in Quadrants scope is not supported. Is and IsNot are refused before the static-scope check, so the natural predicate cannot be written at all	python/quadrants/lang/ast/ast_transformer.py:945
the predicate one has to write instead	qd.static(x), which tests truthiness rather than presence
bias: qd.Tensor | None	rejected with Invalid type annotation (argument 1) of Quadrants kernel	python/quadrants/lang/_func_base.py:281
bias: qd.Tensor = None	rejected with Quadrants kernels do not support default values for arguments, so every call site has to pass None explicitly or wrap the kernel in a thin Python function	python/quadrants/lang/_func_base.py:235
test and documentation coverage	none. No test passes None into a qd.Tensor or qd.template() slot, and neither static.md nor tensor_types.md mentions the behaviour, so it can regress without anything failing

The truthiness fallback is also a trap waiting to be sprung. It is safe for the tensor families, since Ndarray, Field and Tensor define neither __bool__ nor __len__ and a live instance is therefore always truthy, but the same idiom is silently wrong for an integer template argument equal to zero or for an empty container. A reader has no way to tell which case they are looking at.

Finally, the union spelling is exactly the one a type checker wants. qd.Tensor | None is a valid type expression, so accepting it also removes an annotation form that downstream projects currently cannot write, which is the same class of problem as #831.

Describe the solution you'd like

1. Allow Is and IsNot inside a static scope by moving them next to In and NotIn in ops_static in build_Compare, and keep rejecting them in quadrants scope. This alone makes qd.static(x is not None) express the intent, instead of the truthiness workaround.
2. Accept T | None and Optional[T] as kernel parameter annotations for the non-template argument families, that is qd.Tensor and qd.types.NDArray, and route the None value through the template path so the argument specializes on presence. qd.types.NDArray | None matters most here, since it is the only annotation that accepts an external array and it is the one case with no workaround at all today. qd.Template can be accepted in the same form for uniformity, although a template slot already handles None.
3. Allow = None as a default for a parameter whose annotation admits None, so an unset argument can be omitted at the call site.
4. Add tests pinning the specialization count and the absent branch never being traced, plus a short section in static.md, so this becomes a supported contract rather than behaviour that happens to work.

Items 1 and 3 are small and independent of each other. Item 2 is the one that makes the kernel signature read the way it should.

Additional comments

Verified against quadrants 1.1.4 at commit 47e9dfb5, arch=arm64 CPU, with the offline cache disabled. The offline cache path with a None argument slot has not been exercised.

[hughperkins]

Agent said:

Overall I think this is a good issue - unusually well-researched, the code references check out, and the underlying need is real. But I'd restructure the proposal: the mechanism it asks for already exists and works end to end, so the fix is smaller than the issue implies, and I'd build it on the existing qd.Tensor machinery rather than teaching NdarrayType about None.

Everything below is from reading the code at origin/main (ab9a58ab5); nothing was executed, so the runtime claims in the issue itself are taken at face value.

The dual-nature slot already exists

The issue frames qd.Tensor + None as "behaviour that happens to work". It's more deliberate than that. A qd.Tensor slot is simultaneously a compile-time slot and a runtime slot, and which one it is gets decided per specialization:

• It's registered in template_slot_locations (_func_base.py:288), so its Python value is injected into global_vars and is visible to qd.static.
• But the launch loop only skips template-annotated slots (kernel.py:637), so a qd.Tensor slot still reaches _recursive_set_args, which returns 0, True for the non-ndarray branch (_func_base.py:629) and one kernel arg for the ndarray branch.

That last part is the interesting bit. Because the launch loop advances its index with i_out += num_args_, an argument that consumes zero runtime slots in one specialization and one in another already works, and the positional indices of later arguments shift correctly. That is the hard part of "specialize away the absent path", and it is already built and exercised.

So the architectural objection I expected to raise - that the compiled kernel's parameter list would have to vary per specialization, while template_slot_locations is fixed at decoration time - doesn't apply. The design already accommodates it.

I'd do item 2 differently

The issue proposes accepting T | None for both qd.Tensor and qd.types.NDArray, and calls qd.types.NDArray | None the important one because it is the only annotation that accepts a torch tensor.

I'd invert that. The reason qd.types.NDArray matters is only that qd.Tensor rejects external arrays - and that rejection looks incidental rather than intended. In _extract_arg, the qd.Tensor branch tests issubclass(arg_type, (Ndarray, AnyArray)) and routes matches to _TENSOR_T_NDARRAY_ANNOTATION, which is a plain NdarrayType() (_template_mapper_hotpath.py:242-250). The NdarrayType path immediately below it already handles external arrays via the getattr(arg, "shape", None) branch at line 336. So the ndarray feature path that qd.Tensor delegates to can already consume a torch tensor; it is the predicate deciding whether to delegate that is too narrow. Same story at launch, where _recursive_set_args gates on isinstance(v, Ndarray) (_func_base.py:620).

Widening that predicate to include external arrays would fix the torch case, and it would fix the specialization explosion the issue mentions as a side effect: torch tensors would be keyed by dtype and ndim through the ndarray path instead of by object identity through the template path, so two distinct torch tensors would share one specialization rather than compiling two. Then qd.Tensor | None covers every case and NDArray | None doesn't need to exist. One mechanism instead of two parallel ones with subtly different dispatch.

The compatibility risk to check before doing this: anything currently passed to a qd.Tensor slot that has a .shape but is meant as a compile-time template value. A numpy array is the obvious candidate - today it binds as a template constant, and after the change it would become a runtime ndarray argument. Whether that is a real break depends on whether anyone does it.

Items 1 and 3 are nearly free

Item 1 (allowing is / is not in static scope) is a few lines. The rejection at ast_transformer.py:943-945 sits before the ops.get lookup, so it fires unconditionally; it just needs to move inside the non-static branch. Worth noting the current error message says Operator "is not" in Quadrants scope is not supported even when you are demonstrably inside a static scope, which is its own small wart.

Item 3 (allowing = None defaults) is closer to free than the issue suggests. ArgMetadata already stores param.default for kernels (_func_base.py:283), and the argument-fusing path already fills missing arguments from arg_meta.default (line 498), distinguishing "not supplied" via the inspect.Parameter.empty sentinel rather than via None. That path runs on every kernel launch, not just for @qd.func. So the only thing standing in the way is the blanket rejection at line 234. The docstring at line 449 even says "Quadrants doesn't allow defaults, so we don't need to consider default options here" while the code fifty lines down does exactly that.

One caveat on item 3 that the issue doesn't raise: the fuse function has an early-return fast path (if not (kwargs or num_arg_metas > num_args): return py_args). Omitting a trailing argument defeats it and rebuilds a fused_py_args list on every launch. Given how aggressively this file is annotated for per-launch cost - there are comments about a 180 ns pybind penalty and about getattr versus direct attribute access - and given that the motivating Genesis force kernel is a hot one, that is worth measuring, or worth precomputing a defaults tail.

What I'd add to the issue

A few things I think are missing.

Offline cache. It is listed as an unexercised afterthought. I'd promote it to a required test. It is the most likely place for a latent bug, since a None has to survive spec-key serialization and the two specializations have to stay distinct across a cache round-trip.

Autodiff. The ndarray annotation carries needs_grad, and the absent specialization has no gradient slot at all. Worth pinning what happens when an optional tensor kernel is used under autodiff.

Fastcache and purity. The repo's own documentation says a fastcache=True kernel must be pure and must not read globals - but the absent-argument value is delivered through global_vars (function_def_transformer.py:87). So an optional-argument kernel reads a global by construction. Either fastcache already excludes template slots and this is fine, or the two features conflict. Either way it needs a test, because a hot Genesis kernel is exactly the kind of thing someone will want to fastcache.

A design question the issue skips. If qd.Tensor | None becomes the supported spelling, does bare qd.Tensor keep accepting None? Right now it does, silently. Leaving that as-is means the annotation lies about what it accepts, which undercuts the type-checkability argument the issue borrows from #831. Tightening it is a compatibility break. That decision should be made explicitly rather than by default.

Sequencing

Items 1, 3 and 4 are independent and small. I'd land them first as one or two PRs - particularly the tests, which lock in the currently-undocumented behaviour before anyone refactors it away. Item 2 is the real work and deserves its own PR, ideally after settling the qd.Tensor versus NDArray question above, since that determines how much surface area changes.

[duburcqa]

I’m not ok. None support must be consistent across Tensor, ndarray and field annotations. Or ndarray and field annotations should be removed entirely if one is not supposed to use them.

[hughperkins]

Unstructured notes from meeting with Alexis:

qd.Tensor supporting torch makes sense

• not supporting torch arguably a bug currenlty

not sufficient because:

is ndarray just quadrants tensor?

• also external tensor?

ndarray vs field, static vs dynamic sized

1. if ndarray means anyting taht can be interpreted by quadrnats as dyanmic sized array, then should support torch
2. if ndarray and field types are still useful
   vs stop using ndarray and field typing

could ditch ndarray and field types

if keep encouarging / not discouarging direct use of ndarray and field, support of None inconcistency very confusing.

could say, only support qd.Tensor None today, but not to be satisfied by this, open a ticket to track this defect.

supporting None for ndarray and field and lower types is not priority.