upcast array to solve njit failure (#3652)

* upcast array to solve njit failure

* add tests and remove debug statements

shap/utils/_masked_model.py

@@ -356,12 +356,40 @@ def _build_delta_masked_inputs(masks, batch_positions, num_mask_samples, num_var
 
     return all_masked_inputs, i + 1 # i + 1 is the number of output rows after averaging
 
+def _upcast_array(arr: np.ndarray) -> np.ndarray:
+    """Since njit doesn't support float16, we need to upcast it to float32.
+
+    Args:
+        arr (np.ndarray): array to upcast
+
+    Returns
+    -------
+        np.ndarray: upcasted array
+    """
+    if arr.dtype == np.float16:
+        return arr.astype(np.float32)
+    else:
+        return arr
 
 def _build_fixed_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights):
     if len(last_outs.shape) == 1:
-        _build_fixed_single_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights)
+        _build_fixed_single_output(_upcast_array(averaged_outs),
+                                   _upcast_array(last_outs),
+                                   _upcast_array(outputs),
+                                   batch_positions,
+                                   varying_rows,
+                                   num_varying_rows,
+                                   link,
+                                   linearizing_weights)
     else:
-        _build_fixed_multi_output(averaged_outs, last_outs, outputs, batch_positions, varying_rows, num_varying_rows, link, linearizing_weights)
+        _build_fixed_multi_output(_upcast_array(averaged_outs),
+                                  _upcast_array(last_outs),
+                                  _upcast_array(outputs),
+                                  batch_positions,
+                                  varying_rows,
+                                  num_varying_rows,
+                                  link,
+                                  linearizing_weights)
 
 @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):
@@ -389,6 +417,7 @@ def _build_fixed_single_output(averaged_outs, last_outs, outputs, batch_position
 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
+
     sample_count = last_outs.shape[0]
     for i in range(len(averaged_outs)):
         if batch_positions[i] < batch_positions[i+1]:

tests/utils/test_masked_model.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+from shap.links import identity
+from shap.utils._masked_model import _build_fixed_output
+
+
+def test__build_fixed_output():
+    """GH3651"""
+    num_varying_rows = np.array([1])
+    varying_rows = np.array([[True]])
+    batch_positions = np.array([0, 1])
+    averaged_outs = np.zeros((1, 10), dtype=np.float32)
+    last_outs = np.zeros((1, 10), dtype=np.float32)
+    outputs = np.random.rand(1, 10).astype(np.float16)
+    _build_fixed_output(averaged_outs,
+                        last_outs,
+                        outputs,
+                        batch_positions,
+                        varying_rows,
+                        num_varying_rows,
+                        identity,
+                        None)
+    assert np.allclose(averaged_outs, outputs, 1e-2)