Numpy histogram large tensor fix (#19392)

* fix histogram * tweak naming
apache · Oct 22, 2020 · 3f833e1 · 3f833e1
1 parent 43750c8
commit 3f833e1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/operator/tensor/histogram.cc b/src/operator/tensor/histogram.cc
@@ -28,7 +28,7 @@ namespace op {
 
 struct ComputeBinKernel {
   template<typename DType>
-  MSHADOW_XINLINE static void Map(int i, const DType* in_data, const DType* bin_bounds,
+  MSHADOW_XINLINE static void Map(index_t i, const DType* in_data, const DType* bin_bounds,
                                   int* bin_indices, int bin_cnt, double min, double max) {
     DType data = in_data[i];
     int target = -1;
@@ -42,7 +42,7 @@ struct ComputeBinKernel {
   }
 
   template<typename DType>
-  MSHADOW_XINLINE static void Map(int i, const DType* in_data, int* bin_indices,
+  MSHADOW_XINLINE static void Map(index_t i, const DType* in_data, int* bin_indices,
                                    const DType* bin_bounds, int num_bins) {
     DType data = in_data[i];
     int target = -1;

diff --git a/tests/nightly/test_np_large_array.py b/tests/nightly/test_np_large_array.py
@@ -2112,6 +2112,15 @@ def test_dsplit():
     assert inp.grad[-1][-1][0] == 0 and inp.grad[0][1][1] == 1
 
 
+@use_np
+def test_histogram():
+    inp = np.ones((INT_OVERFLOW, 2))
+    inp[-1, -1] = 2
+    hist, _ = np.histogram(inp, np.array([0.5, 1.5, 2.5, 3.5]))
+    assert hist.shape == (3, )
+    assert hist[0] == int(2 * INT_OVERFLOW - 1) and hist[1] == 1
+
+
 @use_np
 def test_nan_to_num():
     inp = np.zeros((3, INT_OVERFLOW))