add inference of unknown dimension value

tgmattso · Apr 21, 2021 · cea2726 · cea2726
1 parent e1112f3
commit cea2726
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/numba/misc/dummyarray.py b/numba/misc/dummyarray.py
@@ -326,6 +326,15 @@ def reshape(self, *newdims, **kws):
             else:
                 knownsize *= dim
 
+        # compute the missing dimension
+        if unknownidx >= 0:
+            if knownsize == 0 or self.size % knownsize != 0:
+                raise ValueError("cannot infer valid shape for unknown dimension")
+            else:
+                newdims = newdims[0:unknownidx] \
+                        + tuple((self.size // knownsize,)) \
+                        + newdims[unknownidx + 1:]
+
         newsize = functools.reduce(operator.mul, newdims, 1)
 
         if order == 'A':