Upload 14 files

monotonic_align/__init__.py

@@ -1,19 +1,19 @@
-import numpy as np
-import torch
-from .monotonic_align.core import maximum_path_c
+from numpy import zeros, int32, float32
+from torch import from_numpy
 
+from .core import maximum_path_jit
 
 def maximum_path(neg_cent, mask):
-  """ Cython optimized version.
+  """ numba optimized version.
   neg_cent: [b, t_t, t_s]
   mask: [b, t_t, t_s]
   """
   device = neg_cent.device
   dtype = neg_cent.dtype
-  neg_cent = neg_cent.data.cpu().numpy().astype(np.float32)
-  path = np.zeros(neg_cent.shape, dtype=np.int32)
+  neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+  path = zeros(neg_cent.shape, dtype=int32)
 
-  t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(np.int32)
-  t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(np.int32)
-  maximum_path_c(path, neg_cent, t_t_max, t_s_max)
-  return torch.from_numpy(path).to(device=device, dtype=dtype)
+  t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+  t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+  maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
+  return from_numpy(path).to(device=device, dtype=dtype)

monotonic_align/core.c

@@ -21,7 +21,7 @@ END: Cython Metadata */
 #else
 #define CYTHON_ABI "0_29_21"
 #define CYTHON_HEX_VERSION 0x001D15F0
-#define CYTHON_FUTURE_DIVISION 0
+#define CYTHON_FUTURE_DIVISION 1
 #include <stddef.h>
 #ifndef offsetof
   #define offsetof(type, member) ( (size_t) & ((type*)0) -> member )
@@ -949,7 +949,7 @@ struct __pyx_memoryview_obj;
 struct __pyx_memoryviewslice_obj;
 struct __pyx_opt_args_15monotonic_align_4core_maximum_path_each;
 
-/* "monotonic_align/core.pyx":7
+/* "monotonic_align/core.pyx":9
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:             # <<<<<<<<<<<<<<
@@ -2031,7 +2031,7 @@ static PyObject *__pyx_tuple__25;
 static PyObject *__pyx_codeobj__26;
 /* Late includes */
 
-/* "monotonic_align/core.pyx":7
+/* "monotonic_align/core.pyx":9
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:             # <<<<<<<<<<<<<<
@@ -2068,7 +2068,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     }
   }
 
-  /* "monotonic_align/core.pyx":13
+  /* "monotonic_align/core.pyx":15
  *   cdef float v_cur
  *   cdef float tmp
  *   cdef int index = t_x - 1             # <<<<<<<<<<<<<<
@@ -2077,7 +2077,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
  */
   __pyx_v_index = (__pyx_v_t_x - 1);
 
-  /* "monotonic_align/core.pyx":15
+  /* "monotonic_align/core.pyx":17
  *   cdef int index = t_x - 1
  * 
  *   for y in range(t_y):             # <<<<<<<<<<<<<<
@@ -2089,7 +2089,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
   for (__pyx_t_3 = 0; __pyx_t_3 < __pyx_t_2; __pyx_t_3+=1) {
     __pyx_v_y = __pyx_t_3;
 
-    /* "monotonic_align/core.pyx":16
+    /* "monotonic_align/core.pyx":18
  * 
  *   for y in range(t_y):
  *     for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):             # <<<<<<<<<<<<<<
@@ -2115,7 +2115,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     for (__pyx_t_5 = __pyx_t_7; __pyx_t_5 < __pyx_t_6; __pyx_t_5+=1) {
       __pyx_v_x = __pyx_t_5;
 
-      /* "monotonic_align/core.pyx":17
+      /* "monotonic_align/core.pyx":19
  *   for y in range(t_y):
  *     for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
  *       if x == y:             # <<<<<<<<<<<<<<
@@ -2125,7 +2125,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
       __pyx_t_8 = ((__pyx_v_x == __pyx_v_y) != 0);
       if (__pyx_t_8) {
 
-        /* "monotonic_align/core.pyx":18
+        /* "monotonic_align/core.pyx":20
  *     for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
  *       if x == y:
  *         v_cur = max_neg_val             # <<<<<<<<<<<<<<
@@ -2134,7 +2134,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
  */
         __pyx_v_v_cur = __pyx_v_max_neg_val;
 
-        /* "monotonic_align/core.pyx":17
+        /* "monotonic_align/core.pyx":19
  *   for y in range(t_y):
  *     for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
  *       if x == y:             # <<<<<<<<<<<<<<
@@ -2144,7 +2144,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
         goto __pyx_L7;
       }
 
-      /* "monotonic_align/core.pyx":20
+      /* "monotonic_align/core.pyx":22
  *         v_cur = max_neg_val
  *       else:
  *         v_cur = value[y-1, x]             # <<<<<<<<<<<<<<
@@ -2158,7 +2158,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
       }
       __pyx_L7:;
 
-      /* "monotonic_align/core.pyx":21
+      /* "monotonic_align/core.pyx":23
  *       else:
  *         v_cur = value[y-1, x]
  *       if x == 0:             # <<<<<<<<<<<<<<
@@ -2168,7 +2168,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
       __pyx_t_8 = ((__pyx_v_x == 0) != 0);
       if (__pyx_t_8) {
 
-        /* "monotonic_align/core.pyx":22
+        /* "monotonic_align/core.pyx":24
  *         v_cur = value[y-1, x]
  *       if x == 0:
  *         if y == 0:             # <<<<<<<<<<<<<<
@@ -2178,7 +2178,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
         __pyx_t_8 = ((__pyx_v_y == 0) != 0);
         if (__pyx_t_8) {
 
-          /* "monotonic_align/core.pyx":23
+          /* "monotonic_align/core.pyx":25
  *       if x == 0:
  *         if y == 0:
  *           v_prev = 0.             # <<<<<<<<<<<<<<
@@ -2187,7 +2187,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
  */
           __pyx_v_v_prev = 0.;
 
-          /* "monotonic_align/core.pyx":22
+          /* "monotonic_align/core.pyx":24
  *         v_cur = value[y-1, x]
  *       if x == 0:
  *         if y == 0:             # <<<<<<<<<<<<<<
@@ -2197,7 +2197,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
           goto __pyx_L9;
         }
 
-        /* "monotonic_align/core.pyx":25
+        /* "monotonic_align/core.pyx":27
  *           v_prev = 0.
  *         else:
  *           v_prev = max_neg_val             # <<<<<<<<<<<<<<
@@ -2209,7 +2209,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
         }
         __pyx_L9:;
 
-        /* "monotonic_align/core.pyx":21
+        /* "monotonic_align/core.pyx":23
  *       else:
  *         v_cur = value[y-1, x]
  *       if x == 0:             # <<<<<<<<<<<<<<
@@ -2219,7 +2219,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
         goto __pyx_L8;
       }
 
-      /* "monotonic_align/core.pyx":27
+      /* "monotonic_align/core.pyx":29
  *           v_prev = max_neg_val
  *       else:
  *         v_prev = value[y-1, x-1]             # <<<<<<<<<<<<<<
@@ -2233,7 +2233,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
       }
       __pyx_L8:;
 
-      /* "monotonic_align/core.pyx":28
+      /* "monotonic_align/core.pyx":30
  *       else:
  *         v_prev = value[y-1, x-1]
  *       value[y, x] += max(v_prev, v_cur)             # <<<<<<<<<<<<<<
@@ -2253,7 +2253,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     }
   }
 
-  /* "monotonic_align/core.pyx":30
+  /* "monotonic_align/core.pyx":32
  *       value[y, x] += max(v_prev, v_cur)
  * 
  *   for y in range(t_y - 1, -1, -1):             # <<<<<<<<<<<<<<
@@ -2263,7 +2263,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
   for (__pyx_t_1 = (__pyx_v_t_y - 1); __pyx_t_1 > -1; __pyx_t_1-=1) {
     __pyx_v_y = __pyx_t_1;
 
-    /* "monotonic_align/core.pyx":31
+    /* "monotonic_align/core.pyx":33
  * 
  *   for y in range(t_y - 1, -1, -1):
  *     path[y, index] = 1             # <<<<<<<<<<<<<<
@@ -2274,7 +2274,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     __pyx_t_9 = __pyx_v_index;
     *((int *) ( /* dim=1 */ ((char *) (((int *) ( /* dim=0 */ (__pyx_v_path.data + __pyx_t_10 * __pyx_v_path.strides[0]) )) + __pyx_t_9)) )) = 1;
 
-    /* "monotonic_align/core.pyx":32
+    /* "monotonic_align/core.pyx":34
  *   for y in range(t_y - 1, -1, -1):
  *     path[y, index] = 1
  *     if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):             # <<<<<<<<<<<<<<
@@ -2302,7 +2302,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     __pyx_L13_bool_binop_done:;
     if (__pyx_t_8) {
 
-      /* "monotonic_align/core.pyx":33
+      /* "monotonic_align/core.pyx":35
  *     path[y, index] = 1
  *     if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):
  *       index = index - 1             # <<<<<<<<<<<<<<
@@ -2311,7 +2311,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
  */
       __pyx_v_index = (__pyx_v_index - 1);
 
-      /* "monotonic_align/core.pyx":32
+      /* "monotonic_align/core.pyx":34
  *   for y in range(t_y - 1, -1, -1):
  *     path[y, index] = 1
  *     if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):             # <<<<<<<<<<<<<<
@@ -2321,7 +2321,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
     }
   }
 
-  /* "monotonic_align/core.pyx":7
+  /* "monotonic_align/core.pyx":9
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:             # <<<<<<<<<<<<<<
@@ -2332,7 +2332,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_each(__Pyx_memviewslice
   /* function exit code */
 }
 
-/* "monotonic_align/core.pyx":38
+/* "monotonic_align/core.pyx":40
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_ys, int[::1] t_xs) nogil:             # <<<<<<<<<<<<<<
@@ -2352,7 +2352,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_c(__Pyx_memviewslice __
   Py_ssize_t __pyx_t_6;
   Py_ssize_t __pyx_t_7;
 
-  /* "monotonic_align/core.pyx":39
+  /* "monotonic_align/core.pyx":41
  * @cython.wraparound(False)
  * cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_ys, int[::1] t_xs) nogil:
  *   cdef int b = paths.shape[0]             # <<<<<<<<<<<<<<
@@ -2361,7 +2361,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_c(__Pyx_memviewslice __
  */
   __pyx_v_b = (__pyx_v_paths.shape[0]);
 
-  /* "monotonic_align/core.pyx":41
+  /* "monotonic_align/core.pyx":43
  *   cdef int b = paths.shape[0]
  *   cdef int i
  *   for i in prange(b, nogil=True):             # <<<<<<<<<<<<<<
@@ -2397,7 +2397,7 @@ static void __pyx_f_15monotonic_align_4core_maximum_path_c(__Pyx_memviewslice __
                         {
                             __pyx_v_i = (int)(0 + 1 * __pyx_t_2);
 
-                            /* "monotonic_align/core.pyx":42
+                            /* "monotonic_align/core.pyx":44
  *   cdef int i
  *   for i in prange(b, nogil=True):
  *     maximum_path_each(paths[i], values[i], t_ys[i], t_xs[i])             # <<<<<<<<<<<<<<
@@ -2458,7 +2458,7 @@ __pyx_t_6 = __pyx_v_i;
         #endif
       }
 
-      /* "monotonic_align/core.pyx":41
+      /* "monotonic_align/core.pyx":43
  *   cdef int b = paths.shape[0]
  *   cdef int i
  *   for i in prange(b, nogil=True):             # <<<<<<<<<<<<<<
@@ -2476,7 +2476,7 @@ __pyx_t_6 = __pyx_v_i;
       }
   }
 
-  /* "monotonic_align/core.pyx":38
+  /* "monotonic_align/core.pyx":40
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_ys, int[::1] t_xs) nogil:             # <<<<<<<<<<<<<<
@@ -2527,23 +2527,23 @@ static PyObject *__pyx_pw_15monotonic_align_4core_1maximum_path_c(PyObject *__py
         case  1:
         if (likely((values[1] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_values)) != 0)) kw_args--;
         else {
-          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 1); __PYX_ERR(0, 38, __pyx_L3_error)
+          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 1); __PYX_ERR(0, 40, __pyx_L3_error)
         }
         CYTHON_FALLTHROUGH;
         case  2:
         if (likely((values[2] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_t_ys)) != 0)) kw_args--;
         else {
-          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 2); __PYX_ERR(0, 38, __pyx_L3_error)
+          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 2); __PYX_ERR(0, 40, __pyx_L3_error)
         }
         CYTHON_FALLTHROUGH;
         case  3:
         if (likely((values[3] = __Pyx_PyDict_GetItemStr(__pyx_kwds, __pyx_n_s_t_xs)) != 0)) kw_args--;
         else {
-          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 3); __PYX_ERR(0, 38, __pyx_L3_error)
+          __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, 3); __PYX_ERR(0, 40, __pyx_L3_error)
         }
       }
       if (unlikely(kw_args > 0)) {
-        if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "maximum_path_c") < 0)) __PYX_ERR(0, 38, __pyx_L3_error)
+        if (unlikely(__Pyx_ParseOptionalKeywords(__pyx_kwds, __pyx_pyargnames, 0, values, pos_args, "maximum_path_c") < 0)) __PYX_ERR(0, 40, __pyx_L3_error)
       }
     } else if (PyTuple_GET_SIZE(__pyx_args) != 4) {
       goto __pyx_L5_argtuple_error;
@@ -2553,14 +2553,14 @@ static PyObject *__pyx_pw_15monotonic_align_4core_1maximum_path_c(PyObject *__py
       values[2] = PyTuple_GET_ITEM(__pyx_args, 2);
       values[3] = PyTuple_GET_ITEM(__pyx_args, 3);
     }
-    __pyx_v_paths = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_paths.memview)) __PYX_ERR(0, 38, __pyx_L3_error)
-    __pyx_v_values = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_values.memview)) __PYX_ERR(0, 38, __pyx_L3_error)
-    __pyx_v_t_ys = __Pyx_PyObject_to_MemoryviewSlice_dc_int(values[2], PyBUF_WRITABLE); if (unlikely(!__pyx_v_t_ys.memview)) __PYX_ERR(0, 38, __pyx_L3_error)
-    __pyx_v_t_xs = __Pyx_PyObject_to_MemoryviewSlice_dc_int(values[3], PyBUF_WRITABLE); if (unlikely(!__pyx_v_t_xs.memview)) __PYX_ERR(0, 38, __pyx_L3_error)
+    __pyx_v_paths = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_int(values[0], PyBUF_WRITABLE); if (unlikely(!__pyx_v_paths.memview)) __PYX_ERR(0, 40, __pyx_L3_error)
+    __pyx_v_values = __Pyx_PyObject_to_MemoryviewSlice_d_d_dc_float(values[1], PyBUF_WRITABLE); if (unlikely(!__pyx_v_values.memview)) __PYX_ERR(0, 40, __pyx_L3_error)
+    __pyx_v_t_ys = __Pyx_PyObject_to_MemoryviewSlice_dc_int(values[2], PyBUF_WRITABLE); if (unlikely(!__pyx_v_t_ys.memview)) __PYX_ERR(0, 40, __pyx_L3_error)
+    __pyx_v_t_xs = __Pyx_PyObject_to_MemoryviewSlice_dc_int(values[3], PyBUF_WRITABLE); if (unlikely(!__pyx_v_t_xs.memview)) __PYX_ERR(0, 40, __pyx_L3_error)
   }
   goto __pyx_L4_argument_unpacking_done;
   __pyx_L5_argtuple_error:;
-  __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 38, __pyx_L3_error)
+  __Pyx_RaiseArgtupleInvalid("maximum_path_c", 1, 4, 4, PyTuple_GET_SIZE(__pyx_args)); __PYX_ERR(0, 40, __pyx_L3_error)
   __pyx_L3_error:;
   __Pyx_AddTraceback("monotonic_align.core.maximum_path_c", __pyx_clineno, __pyx_lineno, __pyx_filename);
   __Pyx_RefNannyFinishContext();
@@ -2582,11 +2582,11 @@ static PyObject *__pyx_pf_15monotonic_align_4core_maximum_path_c(CYTHON_UNUSED P
   int __pyx_clineno = 0;
   __Pyx_RefNannySetupContext("maximum_path_c", 0);
   __Pyx_XDECREF(__pyx_r);
-  if (unlikely(!__pyx_v_paths.memview)) { __Pyx_RaiseUnboundLocalError("paths"); __PYX_ERR(0, 38, __pyx_L1_error) }
-  if (unlikely(!__pyx_v_values.memview)) { __Pyx_RaiseUnboundLocalError("values"); __PYX_ERR(0, 38, __pyx_L1_error) }
-  if (unlikely(!__pyx_v_t_ys.memview)) { __Pyx_RaiseUnboundLocalError("t_ys"); __PYX_ERR(0, 38, __pyx_L1_error) }
-  if (unlikely(!__pyx_v_t_xs.memview)) { __Pyx_RaiseUnboundLocalError("t_xs"); __PYX_ERR(0, 38, __pyx_L1_error) }
-  __pyx_t_1 = __Pyx_void_to_None(__pyx_f_15monotonic_align_4core_maximum_path_c(__pyx_v_paths, __pyx_v_values, __pyx_v_t_ys, __pyx_v_t_xs, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 38, __pyx_L1_error)
+  if (unlikely(!__pyx_v_paths.memview)) { __Pyx_RaiseUnboundLocalError("paths"); __PYX_ERR(0, 40, __pyx_L1_error) }
+  if (unlikely(!__pyx_v_values.memview)) { __Pyx_RaiseUnboundLocalError("values"); __PYX_ERR(0, 40, __pyx_L1_error) }
+  if (unlikely(!__pyx_v_t_ys.memview)) { __Pyx_RaiseUnboundLocalError("t_ys"); __PYX_ERR(0, 40, __pyx_L1_error) }
+  if (unlikely(!__pyx_v_t_xs.memview)) { __Pyx_RaiseUnboundLocalError("t_xs"); __PYX_ERR(0, 40, __pyx_L1_error) }
+  __pyx_t_1 = __Pyx_void_to_None(__pyx_f_15monotonic_align_4core_maximum_path_c(__pyx_v_paths, __pyx_v_values, __pyx_v_t_ys, __pyx_v_t_xs, 0)); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 40, __pyx_L1_error)
   __Pyx_GOTREF(__pyx_t_1);
   __pyx_r = __pyx_t_1;
   __pyx_t_1 = 0;
@@ -16480,7 +16480,7 @@ static __Pyx_StringTabEntry __pyx_string_tab[] = {
   {0, 0, 0, 0, 0, 0, 0}
 };
 static CYTHON_SMALL_CODE int __Pyx_InitCachedBuiltins(void) {
-  __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 15, __pyx_L1_error)
+  __pyx_builtin_range = __Pyx_GetBuiltinName(__pyx_n_s_range); if (!__pyx_builtin_range) __PYX_ERR(0, 17, __pyx_L1_error)
   __pyx_builtin_ValueError = __Pyx_GetBuiltinName(__pyx_n_s_ValueError); if (!__pyx_builtin_ValueError) __PYX_ERR(1, 133, __pyx_L1_error)
   __pyx_builtin_MemoryError = __Pyx_GetBuiltinName(__pyx_n_s_MemoryError); if (!__pyx_builtin_MemoryError) __PYX_ERR(1, 148, __pyx_L1_error)
   __pyx_builtin_enumerate = __Pyx_GetBuiltinName(__pyx_n_s_enumerate); if (!__pyx_builtin_enumerate) __PYX_ERR(1, 151, __pyx_L1_error)
@@ -17107,7 +17107,7 @@ if (!__Pyx_RefNanny) {
   if (__Pyx_patch_abc() < 0) __PYX_ERR(0, 1, __pyx_L1_error)
   #endif
 
-  /* "monotonic_align/core.pyx":7
+  /* "monotonic_align/core.pyx":9
  * @cython.boundscheck(False)
  * @cython.wraparound(False)
  * cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:             # <<<<<<<<<<<<<<
@@ -17117,9 +17117,9 @@ if (!__Pyx_RefNanny) {
   __pyx_k_ = (-1e9);
 
   /* "monotonic_align/core.pyx":1
- * cimport cython             # <<<<<<<<<<<<<<
- * from cython.parallel import prange
+ * # cython:language_level=3             # <<<<<<<<<<<<<<
  * 
+ * cimport cython
  */
   __pyx_t_1 = __Pyx_PyDict_NewPresized(0); if (unlikely(!__pyx_t_1)) __PYX_ERR(0, 1, __pyx_L1_error)
   __Pyx_GOTREF(__pyx_t_1);

monotonic_align/core.py

@@ -0,0 +1,35 @@
+import numba
+
+
+@numba.jit(numba.void(numba.int32[:,:,::1], numba.float32[:,:,::1], numba.int32[::1], numba.int32[::1]), nopython=True, nogil=True)
+def maximum_path_jit(paths, values, t_ys, t_xs):
+  b = paths.shape[0]
+  max_neg_val=-1e9
+  for i in range(int(b)):
+    path = paths[i]
+    value = values[i]
+    t_y = t_ys[i]
+    t_x = t_xs[i]
+
+    v_prev = v_cur = 0.0
+    index = t_x - 1
+
+    for y in range(t_y):
+      for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+        if x == y:
+          v_cur = max_neg_val
+        else:
+          v_cur = value[y-1, x]
+        if x == 0:
+          if y == 0:
+            v_prev = 0.
+          else:
+            v_prev = max_neg_val
+        else:
+          v_prev = value[y-1, x-1]
+        value[y, x] += max(v_prev, v_cur)
+
+    for y in range(t_y - 1, -1, -1):
+      path[y, index] = 1
+      if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):
+        index = index - 1

monotonic_align/core.pyx

@@ -1,3 +1,5 @@
+# cython:language_level=3
+
 cimport cython
 from cython.parallel import prange
 
@@ -39,4 +41,4 @@ cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_
   cdef int b = paths.shape[0]
   cdef int i
   for i in prange(b, nogil=True):
-    maximum_path_each(paths[i], values[i], t_ys[i], t_xs[i])
+    maximum_path_each(paths[i], values[i], t_ys[i], t_xs[i])

monotonic_align/setup.py

@@ -1,3 +1,4 @@
+# cython:language_level=3
 from distutils.core import setup
 from Cython.Build import cythonize
 import numpy
@@ -6,4 +7,4 @@ setup(
   name = 'monotonic_align',
   ext_modules = cythonize("core.pyx"),
   include_dirs=[numpy.get_include()]
-)
+)