diff --git a/hoi/core/tests/test_mi.py b/hoi/core/tests/test_mi.py
index d6bd9f55..a05ebacb 100644
--- a/hoi/core/tests/test_mi.py
+++ b/hoi/core/tests/test_mi.py
@@ -34,8 +34,8 @@ def test_mi_gc(self, xy, biascorrect, copnorm):
     @pytest.mark.parametrize("xy", [(x1, x2), (j1, j2)])
     def test_mi_bin(self, xy):
         mi_fcn = get_mi(method="binning")
-        x_binned, _ = digitize(xy[0], n_bins=3)
-        y_binned, _ = digitize(xy[1], n_bins=3)
+        x_binned = digitize(xy[0], n_bins=3)
+        y_binned = digitize(xy[1], n_bins=3)
         mi = mi_fcn(x_binned, y_binned)
         assert mi.dtype == np.float32
         assert mi.shape == ()
diff --git a/hoi/utils/stats.py b/hoi/utils/stats.py
index ab37e7c8..84ce839c 100644
--- a/hoi/utils/stats.py
+++ b/hoi/utils/stats.py
@@ -181,9 +181,6 @@ def digitize(x, n_bins, axis=0, use_sklearn=False, **kwargs):
         return np.apply_along_axis(digitize_sklearn, axis, x, **kwargs)
 
 
-partial(jax.jit, static_argnums=(1, 2))
-
-
 def digitize_1d_hist(x, n_bins):
     """One dimensional digitization."""
     assert x.ndim == 1