Skip to content

[BUG]: op annotated -> Any fails in a TransformIterator (annotation lowered to pyobject) #10690

Description

@shwina

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Status
    Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions