numba / numba

NumPy aware dynamic Python compiler using LLVM
https://numba.pydata.org/
BSD 2-Clause "Simplified" License
9.6k stars 1.11k forks source link

nb.optional doesn't work with nb.types.NamedTuple #9607

Open ocmob opened 2 weeks ago

ocmob commented 2 weeks ago

It seems like nb.optional doesn't work with nb.types.NamedTuple (it does with nb.types.Tuple, though). Here is the minimum repro:

import numba as nb
from typing import NamedTuple

class Params(NamedTuple):
    param1: float
    param2: int

ParamsNumbaType = nb.types.NamedTuple((nb.f8, nb.i8), Params)

@nb.njit(ParamsNumbaType(ParamsNumbaType))
def func(my_tup: Params):
    return my_tup

tup = Params(3., 3)
func(tup)  # Works

class ParamsOptional(NamedTuple):
    param1: float
    param2: int | None = None

ParamsOptionalNumbaType = nb.types.NamedTuple((nb.f8, nb.optional(nb.i8)), ParamsOptional)

@nb.njit(ParamsOptionalNumbaType(ParamsOptionalNumbaType))
def func2(my_tup: ParamsOptional):
    return my_tup

tup = ParamsOptional(3., 3)
func2(tup)  # Fails
esc commented 2 weeks ago

@ocmob thank you for reporting this. I have updated the reproducer to include the import numba as nb.

esc commented 2 weeks ago

The error message is

 💣 zsh» python issue_9607.py
Traceback (most recent call last):
  File "/Users/esc/git/numba/issue_9607.py", line 28, in <module>
    func2(tup)  # Fails
    ^^^^^^^^^^
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 658, in _explain_matching_error
    raise TypeError(msg)
TypeError: No matching definition for argument type(s) ParamsOptional(float64, int64)
esc commented 2 weeks ago

@ocmob looks like I can reproduce this. I am provisionally labelling this as a typing bug until a second maintainer can verify it, thank you for reporting.

sklam commented 2 weeks ago

I'd classify this as a feature request to expand support to typing.NamedTuple. All tuple/namedtuple support predates that.

sklam commented 2 weeks ago

Btw, the problem is happening at unboxing. Numba needs to read into the annotations in ParamsOptional to recognize that the second field should be a optional type.