Fix dtype conversion for scalar Numba Dot results

aesara-devs · May 12, 2021 · 9926e07 · 9926e07
1 parent 54e19a6
commit 9926e07
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/aesara/link/numba/dispatch.py b/aesara/link/numba/dispatch.py
@@ -1424,7 +1424,7 @@ def numba_funcify_Dot(op, node, **kwargs):
 
     @numba.njit
     def dot(x, y):
-        return np.dot(inputs_cast(x), inputs_cast(y)).astype(out_dtype)
+        return np.asarray(np.dot(inputs_cast(x), inputs_cast(y))).astype(out_dtype)
 
     return dot
 

diff --git a/tests/link/test_numba.py b/tests/link/test_numba.py
@@ -1694,6 +1694,11 @@ def test_BroadcastTo(x, shape, exc):
             ),
             None,
         ),
+        (
+            set_test_value(aet.lvector(), np.random.random(size=(2,)).astype(np.int64)),
+            set_test_value(aet.lvector(), np.random.random(size=(2,)).astype(np.int64)),
+            None,
+        ),
     ],
 )
 def test_Dot(x, y, exc):
-Original file line number
+Diff line change
@@ Expand Up / @@ -1694,6 +1694,11 @@ def test_BroadcastTo(x, shape, exc): @@
                 ),
                 None,
             ),
+            (
+                set_test_value(aet.lvector(), np.random.random(size=(2,)).astype(np.int64)),
+                set_test_value(aet.lvector(), np.random.random(size=(2,)).astype(np.int64)),
+                None,
+            ),
         ],
     )
     def test_Dot(x, y, exc):
@@ Expand Down @@