[WIP] Generic jitclass

[AlexanderKalistratov]

This is WIP PR for generic jitclasses.
At the moment everything is in a single file, but is going to be split and tests to be added.

[luk-f-a]

hi @AlexanderKalistratov , as I said in the call, this is a great idea, I'm very happy that someone is doing this. Thanks a lot!

I've done some work with jitclass syntax before and I have a (private and ugly) jitclass implementation based on named tuples. I hope I can provide some useful feedback based on my experience. I will do it in stages, so I'll write what I see now and I will add more comments later.

The first thing that comes to mind is that I wouldn't make automatic translations of python types into numba types. I think this will lead to user confusion. I also think this will lead to mypy errors. For example:

@numba.experimental.jitclass({'capacity': ListType(float64)})
class Foo(object):
    def __init__(self):
        self.bar = TList.empty_list(float64)

cap: List = Foo().bar

This does not fail today, because TypedList has not been annotated, so mypy reads it as Any. However, after numba/numba#5958, I think this will fail, as mypy will recognize Foo().bar as typed.List which will not match List. I think (although I'm not 100% sure) that the same will happen if we write

@numba.experimental.jitclass
class Foo(object):
     bar: List[float]

    def __init__(self):
        self.bar = TList.empty_list(float64)

cap: List = Foo().bar

[luk-f-a]

Since this new jitclass will probably have to co-exist with the existing jitclass for a while, I recommend that the name is changed to something else (jitclass2? newjitclass?) for now. Once it's deemed ready to replace the jitclass, then the name can be switched.

[luk-f-a]

I tried using the same example in main.py in another file (tests.py). Somehow I got this error

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Internal error at <numba.core.typeinfer.CallConstraint object at 0x7fe797e0be20>.
Failed in nopython mode pipeline (step: nopython mode backend)
Can't pickle <class 'numba_extras.jitclass.main.Test'>: attribute lookup Test on numba_extras.jitclass.main failed
During: resolving callee type: type(CPUDispatcher(<function ctor at 0x7fe797e5f940>))
During: typing of call at /home/.../mypyprojects/numba-extras/numba_extras/jitclass/tests.py (62)

[luk-f-a]

On the topic of Annotated, I read the PEP again (https://www.python.org/dev/peps/pep-0593/), I think the second part will be ignored. From the rationale section: "a type T can be annotated with metadata x via the typehint Annotated[T, x]", "If a library (or tool) encounters a typehint Annotated[T, x] and has no special logic for metadata x, it should ignore it and simply treat the type as T", "mypy or Pyre, which can safely ignore x".

My understanding of Annotated is that the metadata (the second argument) will be ignored by mypy and I would expect that it's possible to write Annotated[Any, numba.types.float64[:, :,:]]. For example,

class Foo():
    pass

a: Annotated[float, Foo] = Foo() #  error: Incompatible types in assignment (expression has type "Foo", variable has type "float")
reveal_type(a) # note: Revealed type is 'builtins.float'

a: Annotated[Foo, float] = Foo() # no error
reveal_type(a) # note: Revealed type is 'scratch_1.Foo'
b: float = a # error: Incompatible types in assignment (expression has type "Foo", variable has type "float")

a: Annotated[Foo, "types.float64[:]"] = Foo()
reveal_type(a) # note: Revealed type is 'scratch_1.Foo'

[AlexanderKalistratov]

@luk-f-a Oh, I see. It seems I've misunderstand you on dev meeting. Yes, your proposal looks very interesting

[AlexanderKalistratov]

@luk-f-a but it is not generic. Need to see if we can overcome it somehow

[luk-f-a]

@luk-f-a but it is not generic. Need to see if we can overcome it somehow

I see your point. There's a class types.npytypes.Array that could be used. For example types.npytypes.Array(types.float64, ndim=3, layout='C') is equivalent to float64[:,:,::1].
Maybe if __class_getitem__ were implemented on that class, it could return Annotated[np.ndarray, "something that numba can understand"] as way to keep mypy happy and provide numba enough information for delayed typing the array.

[luk-f-a]

btw, there's a new version of mypy with a useful improvement: python/mypy#9625
Annotated can now take numba type instances, instead of only type classes or strings.