api: fix tensor mul

devitocodes · Jul 11, 2024 · 60ea7b7 · 60ea7b7
1 parent 0124871
commit 60ea7b7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/devito/types/basic.py b/devito/types/basic.py
@@ -773,7 +773,7 @@ def _eval_matrix_mul(self, other):
                 new_mat[i] = sum(vec)
 
         # Get new class and return product
-        newcls = self.classof_prod(other, new_mat)
+        newcls = self.classof_prod(other, other.cols)
         return newcls._new(self.rows, other.cols, new_mat, copy=False)
 
 

diff --git a/devito/types/tensor.py b/devito/types/tensor.py
@@ -257,11 +257,8 @@ def new_from_mat(self, mat):
         func = tens_func(self)
         return func._new(self.rows, self.cols, mat)
 
-    def classof_prod(self, other, mat):
-        try:
-            is_mat = len(mat[0]) > 1
-        except TypeError:
-            is_mat = False
+    def classof_prod(self, other, cols):
+        is_mat = cols > 1
         is_time = (getattr(self, '_is_TimeDependent', False) or
                    getattr(other, '_is_TimeDependent', False))
         return mat_time_dict[(is_time, is_mat)]