fix tabulate_jet

firedrakeproject · Nov 7, 2023 · 34f622c · 34f622c
1 parent 890adb7
commit 34f622c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/FIAT/expansions.py b/FIAT/expansions.py
@@ -9,6 +9,7 @@
 
 import numpy
 import math
+from itertools import product
 from FIAT import reference_element
 from FIAT import jacobi
 
@@ -269,7 +270,7 @@ def tabulate_jet(self, n, pts, order=1):
         data = [v0]
         for r in range(1, order+1):
             v = numpy.zeros((D,)*r + v0.shape, dtype=v0.dtype)
-            for index in zip(*[range(D) for k in range(r)]):
+            for index in product(range(D), repeat=r):
                 v[index] = vals[tuple(map(index.count, range(D)))]
             data.append(v.transpose((r, r+1) + tuple(range(r))))
         return data