Repository: numba/numba

Issue:
Problem lowering code using timedelta arrays
Certain numba implementations of numpy functions rely on the `type` function. For example: 

https://github.com/numba/numba/blob/60d2bddd0a9ddf84cf900a70e4b7e17bcbe5dc87/numba/targets/arraymath.py#L246

However, `type` does not work with `timedelta64`, which means that for example, `np.sum(array, axis)` does not work with `timedelta64`.

```python
@numba.njit
def foo(x):
    return np.sum(x,0)

foo(np.ones((2,2))) #works
baz(np.ones((2,2)).astype(np.timedelta64)) #fails
```

This comes from the failure to run `type` on `timedelta64`

```python
def foo(x):
    return 1


@numba.extending.overload(foo)
def foo_fc(x):
    zero = getattr(x, 'dtype', x)(0)
    def foo_impl(x):
        a = type(zero)
        return x
    return foo_impl
    
@numba.njit
def bar(x):
    return foo(x)

bar(np.ones(2)) #works
bar(np.ones(2).astype(np.timedelta64)) #fails
```

Discussion:

Patch:
diff --git a/numba/core/typing/builtins.py b/numba/core/typing/builtins.py
index 0aa397137..3d5dabb8a 100644
--- a/numba/core/typing/builtins.py
+++ b/numba/core/typing/builtins.py
@@ -704,6 +704,22 @@ class NumberAttribute(AttributeTemplate):
             return signature(ty)
 
 
+@infer_getattr
+class NPTimedeltaAttribute(AttributeTemplate):
+    key = types.NPTimedelta
+
+    def resolve___class__(self, ty):
+        return types.NumberClass(ty)
+
+
+@infer_getattr
+class NPDatetimeAttribute(AttributeTemplate):
+    key = types.NPDatetime
+
+    def resolve___class__(self, ty):
+        return types.NumberClass(ty)
+
+
 @infer_getattr
 class SliceAttribute(AttributeTemplate):
     key = types.SliceType
diff --git a/numba/core/typing/npydecl.py b/numba/core/typing/npydecl.py
index 395bfedaa..854d979ee 100644
--- a/numba/core/typing/npydecl.py
+++ b/numba/core/typing/npydecl.py
@@ -125,8 +125,8 @@ class Numpy_rules_ufunc(AbstractTemplate):
         out = list(explicit_outputs)
         implicit_output_count = ufunc.nout - len(explicit_outputs)
         if implicit_output_count > 0:
-            # XXX this is currently wrong for datetime64 and timedelta64,
-            # as ufunc_find_matching_loop() doesn't do any type inference.
+            # XXX this is sometimes wrong for datetime64 and timedelta64,
+            # as ufunc_find_matching_loop() doesn't do any type inference
             ret_tys = ufunc_loop.outputs[-implicit_output_count:]
             if ndims > 0:
                 assert layout is not None
diff --git a/numba/np/npyimpl.py b/numba/np/npyimpl.py
index 9e4473096..3e9547ec4 100644
--- a/numba/np/npyimpl.py
+++ b/numba/np/npyimpl.py
@@ -175,7 +175,8 @@ def _prepare_argument(ctxt, bld, inp, tyinp, where='input operand'):
         strides = cgutils.unpack_tuple(bld, ary.strides, tyinp.ndim)
         return _ArrayHelper(ctxt, bld, shape, strides, ary.data,
                             tyinp.layout, tyinp.dtype, tyinp.ndim, inp)
-    elif types.unliteral(tyinp) in types.number_domain | set([types.boolean]):
+    elif (types.unliteral(tyinp) in types.number_domain | {types.boolean}
+          or isinstance(tyinp, types.NPTimedelta)):
         return _ScalarHelper(ctxt, bld, inp, tyinp)
     else:
         raise NotImplementedError('unsupported type for {0}: {1}'.format(where, str(tyinp)))
diff --git a/numba/np/numpy_support.py b/numba/np/numpy_support.py
index b4cfdbd9d..8ea43b5dc 100644
--- a/numba/np/numpy_support.py
+++ b/numba/np/numpy_support.py
@@ -382,6 +382,35 @@ def ufunc_find_matching_loop(ufunc, arg_types):
                   for letter in ufunc_letters[len(numba_types):]]
         return types
 
+    def set_output_dt_units(inputs, outputs):
+        """
+        Sets the output unit of a datetime type based on the input units
+
+        Timedelta is a special dtype that requires the time unit to be
+        specified (day, month, etc). Not every operation with timedelta inputs
+        leads to an output of timedelta output. However, for those that do,
+        the unit of output must be inferred based on the units of the inputs.
+
+        At the moment this function takes care of the case where all
+        inputs have the same unit, and therefore the output unit has the same.
+        If in the future this should be extended to a case with mixed units,
+        the rules should be implemented in `npdatetime_helpers` and called
+        from this function to set the correct output unit.
+        """
+        if all(inp.unit == inputs[0].unit for inp in inputs):
+            # Case with operation on same units. Operations on different units
+            # not adjusted for now but might need to be added in the future
+            unit = inputs[0].unit
+            new_outputs = []
+            for out in outputs:
+                if isinstance(out, types.NPTimedelta) and out.unit == "":
+                    new_outputs.append(types.NPTimedelta(unit))
+                else:
+                    new_outputs.append(out)
+        else:
+            return outputs
+        return new_outputs
+
     # In NumPy, the loops are evaluated from first to last. The first one
     # that is viable is the one used. One loop is viable if it is possible
     # to cast every input operand to the one expected by the ufunc.
@@ -423,6 +452,9 @@ def ufunc_find_matching_loop(ufunc, arg_types):
             try:
                 inputs = choose_types(input_types, ufunc_inputs)
                 outputs = choose_types(output_types, ufunc_outputs)
+                if ufunc_inputs == 'mm':
+                    outputs = set_output_dt_units(inputs, outputs)
+
             except NotImplementedError:
                 # One of the selected dtypes isn't supported by Numba
                 # (e.g. float16), try other candidates


Test patch:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
index 685f02c7a..f9c910a79 100644
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -13,6 +13,9 @@ from numba.tests.support import (TestCase, CompilationCache, MemoryLeak,
                                  MemoryLeakMixin, tag, needs_blas)
 import unittest
 
+TIMEDELTA_M = 'timedelta64[M]'
+TIMEDELTA_Y = 'timedelta64[Y]'
+
 
 def np_around_array(arr, decimals, out):
     np.around(arr, decimals, out)
@@ -980,15 +983,12 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
         pyfunc = array_sum_axis_kws
         cfunc = jit(nopython=True)(pyfunc)
         all_dtypes = [np.float64, np.float32, np.int64, np.uint64, np.complex64,
-                      np.complex128]
-        # timedelta test cannot be enabled until issue #4540 is fixed
-        # all_dtypes += [np.timedelta64]
+                      np.complex128, TIMEDELTA_M]
         all_test_arrays = [
             [np.ones((7, 6, 5, 4, 3), arr_dtype),
              np.ones(1, arr_dtype),
              np.ones((7, 3), arr_dtype) * -5]
             for arr_dtype in all_dtypes]
-
         for arr_list in all_test_arrays:
             for arr in arr_list:
                 for axis in (0, 1, 2):
@@ -1010,12 +1010,11 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
         """
         pyfunc = array_sum_axis_kws
         cfunc = jit(nopython=True)(pyfunc)
-        all_dtypes = [np.int32, np.uint32]
+        all_dtypes = [np.int32, np.uint32, ]
         # expected return dtypes in Numba
         out_dtypes = {np.dtype('int32'): np.int64, np.dtype('uint32'): np.uint64,
-                      np.dtype('int64'): np.int64}
-        # timedelta test cannot be enabled until issue #4540 is fixed
-        # all_dtypes += [np.timedelta64]
+                      np.dtype('int64'): np.int64,
+                      np.dtype(TIMEDELTA_M): np.dtype(TIMEDELTA_M)}
         all_test_arrays = [
             [np.ones((7, 6, 5, 4, 3), arr_dtype),
              np.ones(1, arr_dtype),
@@ -1044,9 +1043,7 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
         pyfunc = array_sum_dtype_kws
         cfunc = jit(nopython=True)(pyfunc)
         all_dtypes = [np.float64, np.float32, np.int64, np.int32, np.uint32,
-                      np.uint64, np.complex64, np.complex128]
-        # timedelta test cannot be enabled until issue #4540 is fixed
-        # all_dtypes += [np.timedelta64]
+                      np.uint64, np.complex64, np.complex128, TIMEDELTA_M]
         all_test_arrays = [
             [np.ones((7, 6, 5, 4, 3), arr_dtype),
              np.ones(1, arr_dtype),
@@ -1061,7 +1058,7 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
                       np.dtype('uint64'): [np.float64, np.int64],
                       np.dtype('complex64'): [np.complex64, np.complex128],
                       np.dtype('complex128'): [np.complex128],
-                      np.dtype('timedelta64'): [np.timedelta64]}
+                      np.dtype(TIMEDELTA_M): [np.dtype(TIMEDELTA_M)]}
 
         for arr_list in all_test_arrays:
             for arr in arr_list:
@@ -1077,7 +1074,7 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
         pyfunc = array_sum_axis_dtype_kws
         cfunc = jit(nopython=True)(pyfunc)
         all_dtypes = [np.float64, np.float32, np.int64, np.int32, np.uint32,
-                      np.uint64, np.complex64, np.complex128]
+                      np.uint64, np.complex64, np.complex128, TIMEDELTA_M]
         all_test_arrays = [
             [np.ones((7, 6, 5, 4, 3), arr_dtype),
              np.ones(1, arr_dtype),
@@ -1092,7 +1089,8 @@ class TestArrayMethods(MemoryLeakMixin, TestCase):
                       np.dtype('uint64'): [np.float64, np.uint64],
                       np.dtype('complex64'): [np.complex64, np.complex128],
                       np.dtype('complex128'): [np.complex128],
-                      np.dtype('timedelta64'): [np.timedelta64]}
+                      np.dtype(TIMEDELTA_M): [np.dtype(TIMEDELTA_M)],
+                      np.dtype(TIMEDELTA_Y): [np.dtype(TIMEDELTA_Y)]}
 
         for arr_list in all_test_arrays:
             for arr in arr_list: