Change datetime64 type within a jitted func

Hi amazing team!

Assume we have an array t which is a np.datetime64. I want to operate with him inside an njitted function (change the type of the datetime and do some math ops. afterwards). Is that possible?

So far I’ve got this:

from numba.types import NPDatetime

@jit(nopython=True)
def foo(dt64_arr):
    return dt64_arr.astype(NPDatetime("M"))
foo(t)

Yet the error I got is the following:

Blockquote
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of <class ‘numba.core.types.scalars.NPDatetime’> with parameters (Literalstr)
No type info available for <class ‘numba.core.types.scalars.NPDatetime’> as a callable.
During: resolving callee type: typeref[<class ‘numba.core.types.scalars.NPDatetime’>]
During: typing of call at /tmp/ipykernel_961918/589338060.py (5)
File “…/…/…/…/tmp/ipykernel_961918/589338060.py”, line 5:

Can you supply a complete example that people can test?

1 Like

Thanks a lot for your prompt reply Nelson, sure thing.

I expect numba to behave like this

import numpy as np
t = np.arange('2005-02', '2005-06', dtype='datetime64[D]')

def foo_np(dt64_arr):
    return dt64_arr.astype("datetime64[M]")
foo_np(t)

To do so, i used the function described above

from numba.types import NPDatetime
t = np.arange('2005-02', '2005-06', dtype='datetime64[D]')

@jit(nopython=True)
def foo_nb(dt64_arr):
    return dt64_arr.astype(NPDatetime("M"))
foo_nb(t)

But it returns an error, described in the first post.
My final goal is to do something like this:

  def foo_final(self):
 
      return self.astype('datetime64[M]').astype(np.int32) % 12 + 1

But using an independent numba function

Okay, I’ll make some comments. I don’t know to use numpy.astype() to convert generic parameterized types like datetime64 from ‘D’ to ‘M’.

That being said, a couple of observations:

  1. if you can do the units-conversion in non-jitted code, or an objmode block, you can use the pointer to the numpy data to do the cast to int32 and other math directly in a jitted function
  2. if you want to/are willing to implement the D-to-M conversion yourself, getting the integer values and doing that logic in the numba block seems possible. There’s a numpy reference here

Either/both of these options are workarounds that may be more work than you want to do for your problem. Perhaps someone else on the board will have better options :slight_smile: