NVIDIA / NVIDIA/cccl

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

Open
#10,690 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

## 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
# 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

```python
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

Contributor guide

Open the contributing guide

Research direction

Read python/cuda_cccl/cuda/compute/iterators/_transform.py around line 85 and run the supplied reproducer to trace how signature_from_annotations handles Any before return-type inference. Compare the TransformIterator and unary_transform paths. Done means the Any case no longer produces the pyobject conversion failure, or reports the proposed actionable annotation error.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.