Bug
numba reads an op's type annotations even under from __future__ import annotations. typing.Any lowers to a pyobject, so an op annotated -> Any fails type inference:
TypingError: No conversion from float64 to pyobject for '$12return_value.3'
The traceback points at the op body and says nothing about annotations.
It is also inconsistent: the same op works with unary_transform, where d_out supplies the return type, and fails only where the type must be inferred. TransformIterator prefers annotations over inference:
# python/cuda_cccl/cuda/compute/iterators/_transform.py:85
_, value_type = signature_from_annotations(transform_op)
if value_type is None:
value_type = self._transform_op.get_return_type((underlying.value_type,))
Reproducer
from __future__ import annotations
from typing import Any
import numpy as np, cupy as cp
import cuda.compute as cc
def square_any(x: Any) -> Any: return x * x
def square_bare(x): return x * x
data = cp.arange(8, dtype=cp.float64); h = np.zeros(1, np.float64)
for tag, op in (("x: Any -> Any", square_any), ("unannotated ", square_bare)):
out = cp.empty(1, cp.float64)
try:
cc.reduce_into(d_in=cc.TransformIterator(data, op), d_out=out,
op=cc.OpKind.PLUS, num_items=8, h_init=h)
print(f"TransformIterator, {tag}: OK ({out.get()[0]})")
except Exception as e:
print(f"TransformIterator, {tag}: FAIL {type(e).__name__}: {str(e).splitlines()[2]}")
cc.unary_transform(d_in=data, d_out=cp.empty(8, cp.float64), op=op, num_items=8)
print(f"unary_transform, {tag}: OK")
Output:
TransformIterator, x: Any -> Any: FAIL TypingError: No conversion from float64 to pyobject for '$12return_value.3'
unary_transform, x: Any -> Any: OK
TransformIterator, unannotated : OK (140.0)
unary_transform, unannotated : OK
A concrete annotation (-> float) also works. The workaround is trivial once you know; finding it is the problem.
Impact
Any is what a type-checked downstream reaches for. cupy ships no stubs, and cuda.compute has no py.typed marker (cuda/cccl does), so anything touching either is Any under mypy. In a mypy-strict project the annotation that satisfies the type checker is the one that breaks the kernel, with an error that points somewhere else.
Suggested fix
Treat annotations that are not numba-typeable (Any, object, bare TypeVar) as absent and fall through to inference, rather than lowering them to a pyobject. Failing that, detect it in the Python layer:
TypeError: operator 'square' is annotated `-> Any`, which numba lowers to a Python
object. Use a concrete annotation (e.g. `-> float`) or remove the annotation.
Separately, shipping py.typed in cuda/compute would remove one reason users end up with Any, and clear the import-untyped error mypy reports for import cuda.compute.
Environment
cuda-cccl 1.1.1
- CUDA 13.3, NVIDIA RTX 6000 Ada, driver 580.167.08
- Python 3.12, cupy-cuda13x 14.1.1, numba-cuda 0.30.4
Bug
numba reads an op's type annotations even under
from __future__ import annotations.typing.Anylowers to a pyobject, so an op annotated-> Anyfails type inference:The traceback points at the op body and says nothing about annotations.
It is also inconsistent: the same op works with
unary_transform, whered_outsupplies the return type, and fails only where the type must be inferred.TransformIteratorprefers annotations over inference:Reproducer
Output:
A concrete annotation (
-> float) also works. The workaround is trivial once you know; finding it is the problem.Impact
Anyis what a type-checked downstream reaches for.cupyships no stubs, andcuda.computehas nopy.typedmarker (cuda/cccldoes), so anything touching either isAnyunder mypy. In a mypy-strict project the annotation that satisfies the type checker is the one that breaks the kernel, with an error that points somewhere else.Suggested fix
Treat annotations that are not numba-typeable (
Any,object, bareTypeVar) as absent and fall through to inference, rather than lowering them to a pyobject. Failing that, detect it in the Python layer:Separately, shipping
py.typedincuda/computewould remove one reason users end up withAny, and clear theimport-untypederror mypy reports forimport cuda.compute.Environment
cuda-cccl1.1.1