migrate from jit to njit

fasttreeshap/links.py

@@ -1,22 +1,22 @@
 import numpy as np
-import numba
+from numba import njit
 
-@numba.jit
+@njit
 def identity(x):
     """ A no-op link function.
     """
     return x
-@numba.jit
+@njit
 def _identity_inverse(x):
     return x
 identity.inverse = _identity_inverse
 
-@numba.jit
+@njit
 def logit(x):
     """ A logit link function useful for going from probability units to log-odds units.
     """
     return np.log(x/(1-x))
-@numba.jit
+@njit
 def _logit_inverse(x):
     return 1/(1+np.exp(-x))
 logit.inverse = _logit_inverse

fasttreeshap/maskers/_tabular.py

@@ -1,7 +1,7 @@
 import logging
 import pandas as pd
 import numpy as np
-from numba import jit
+from numba import njit
 from .. import utils
 from ..utils import safe_isinstance, MaskedModel
 from ._masker import Masker
@@ -181,7 +181,7 @@ def load(cls, in_file, instantiate=True):
             kwargs["clustering"] = s.load("clustering")
         return kwargs
 
-@jit
+@njit
 def _single_delta_mask(dind, masked_inputs, last_mask, data, x, noop_code):
     if dind == noop_code:
         pass
@@ -192,7 +192,7 @@ def _single_delta_mask(dind, masked_inputs, last_mask, data, x, noop_code):
         masked_inputs[:, dind] = x[dind]
         last_mask[dind] = True
 
-@jit
+@njit
 def _delta_masking(masks, x, curr_delta_inds, varying_rows_out,
                    masked_inputs_tmp, last_mask, data, variants, masked_inputs_out, noop_code):
     """ Implements the special (high speed) delta masking API that only flips the positions we need to.

fasttreeshap/utils/_clustering.py

@@ -1,7 +1,7 @@
 import numpy as np
 import scipy as sp
 from scipy.spatial.distance import pdist
-from numba import jit
+from numba import njit
 import sklearn
 import warnings
 from ._general import safe_isinstance
@@ -31,7 +31,7 @@ def partition_tree_shuffle(indexes, index_mask, partition_tree):
     M = len(index_mask)
     #switch = np.random.randn(M) < 0
     _pt_shuffle_rec(partition_tree.shape[0]-1, indexes, index_mask, partition_tree, M, 0)
-@jit
+@njit
 def _pt_shuffle_rec(i, indexes, index_mask, partition_tree, M, pos):
     if i < 0:
         # see if we should include this index in the ordering
@@ -50,7 +50,7 @@ def _pt_shuffle_rec(i, indexes, index_mask, partition_tree, M, pos):
         pos = _pt_shuffle_rec(left, indexes, index_mask, partition_tree, M, pos)
     return pos
 
-@jit
+@njit
 def delta_minimization_order(all_masks, max_swap_size=100, num_passes=2):
     order = np.arange(len(all_masks))
     for _ in range(num_passes):
@@ -59,21 +59,21 @@ def delta_minimization_order(all_masks, max_swap_size=100, num_passes=2):
                 if _reverse_window_score_gain(all_masks, order, i, length) > 0:
                     _reverse_window(order, i, length)
     return order
-@jit
+@njit
 def _reverse_window(order, start, length):
     for i in range(length // 2):
         tmp = order[start + i]
         order[start + i] = order[start + length - i - 1]
         order[start + length - i - 1] = tmp
-@jit
+@njit
 def _reverse_window_score_gain(masks, order, start, length):
     forward_score = _mask_delta_score(masks[order[start - 1]], masks[order[start]]) + \
                     _mask_delta_score(masks[order[start + length-1]], masks[order[start + length]])
     reverse_score = _mask_delta_score(masks[order[start - 1]], masks[order[start + length-1]]) + \
                     _mask_delta_score(masks[order[start]], masks[order[start + length]])
 
     return forward_score - reverse_score
-@jit
+@njit
 def _mask_delta_score(m1, m2):
     return (m1 ^ m2).sum()
 

fasttreeshap/utils/_masked_model.py

@@ -1,7 +1,7 @@
 import copy
 import numpy as np
 import scipy.sparse
-from numba import jit
+from numba import njit
 from .. import links
 
 
@@ -287,7 +287,7 @@ def _convert_delta_mask_to_full(masks, full_masks):
             full_masks[i,masks[masks_pos]] = ~full_masks[i,masks[masks_pos]]
         masks_pos += 1
 
-#@jit # TODO: figure out how to jit this function, or most of it
+#@njit # TODO: figure out how to jit this function, or most of it
 def _build_delta_masked_inputs(masks, batch_positions, num_mask_samples, num_varying_rows, delta_indexes,
                                varying_rows, args, masker, variants, variants_column_sums):
     all_masked_inputs = [[] for a in args]
@@ -358,7 +358,7 @@ def _build_fixed_output(averaged_outs, last_outs, outputs, batch_positions, vary
     else:
         _build_fixed_multi_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights)
 
-@jit # we can't use this when using a custom link function...
+@njit # we can't use this when using a custom link function...
 def _build_fixed_single_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights):
     # here we can assume that the outputs will always be the same size, and we need
     # to carry over evaluation outputs
@@ -380,7 +380,7 @@ def _build_fixed_single_output(averaged_outs, last_outs, outputs, batch_position
         else:
             averaged_outs[i] = averaged_outs[i-1]
 
-@jit
+@njit
 def _build_fixed_multi_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights):
     # here we can assume that the outputs will always be the same size, and we need
     # to carry over evaluation outputs