Closes Bears-R-Us#3870: bug in reshape for bigint type

ajpotts · Nov 20, 2024 · 1f29f03 · 1f29f03
1 parent b456d3e
commit 1f29f03
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/src/AryUtil.chpl b/src/AryUtil.chpl
@@ -15,6 +15,7 @@ module AryUtil
     use List;
     use CommAggregation;
     use CommPrimitives;
+    use BigInteger;
 
 
     param bitsPerDigit = RSLSD_bitsPerDigit;
@@ -905,7 +906,8 @@ module AryUtil
     /*
       unflatten a 1D array into a multi-dimensional array of the given shape
     */
-    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws {
+    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws 
+      where t!=bigint {
       var unflat = makeDistArray((...shape), t);
 
       if N == 1 {
@@ -952,7 +954,6 @@ module AryUtil
             // flat region is spread across multiple locales, do a get for each source locale
             for locInID in locInStart..locInStop {
               const flatSubSlice = flatSlice[flatLocRanges[locInID]];
-
               get(
                 c_ptrTo(unflat[dufc.orderToIndex(flatSubSlice.low)]),
                 getAddr(a[flatSubSlice.low]),
@@ -967,6 +968,24 @@ module AryUtil
       return unflat;
     }
 
+    proc unflatten(const ref a: [?d] ?t, shape: ?N*int): [] t throws 
+      where t==bigint {
+      var unflat = makeDistArray((...shape), t);
+
+      if N == 1 {
+        unflat = a;
+        return unflat;
+      }
+
+      coforall loc in Locales with (ref unflat) do on loc {
+        forall idx in a.localSubdomain() with (var agg = newDstAggregator(t)) {
+                agg.copy(unflat[unflat.domain.orderToIndex(idx)], a[idx]);
+        }
+      }
+
+      return unflat;
+    }
+
     /*
       flatten a multi-dimensional array into a 1D array
     */

diff --git a/tests/pdarrayclass_test.py b/tests/pdarrayclass_test.py
@@ -23,12 +23,20 @@
 class TestPdarrayClass:
 
     @pytest.mark.skip_if_max_rank_less_than(2)
-    def test_reshape(self):
-        a = ak.arange(4)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_reshape(self, dtype):
+        a = ak.arange(4, dtype=dtype)
         r = a.reshape((2, 2))
         assert r.shape == (2, 2)
         assert isinstance(r, ak.pdarray)
 
+    @pytest.mark.skip_if_max_rank_less_than(2)
+    def test_reshape_bug_reproducer(self):
+        dtype = "bigint"
+        size = 10
+        x = ak.arange(size, dtype=dtype).reshape((1, size, 1))
+        ak_assert_equal(x.flatten(), ak.arange(size, dtype=dtype))
+
     def test_shape(self):
         a = ak.arange(4)
         np_a = np.arange(4)