diff --git a/python/dune/perftool/pdelab/tensors.py b/python/dune/perftool/pdelab/tensors.py index 16cf6bb6a4d7bb3c1815ca372f3cdfdd53617167..0b5789ce256b7e6080aa8c0eb1aa1b9bb2ab2d8f 100644 --- a/python/dune/perftool/pdelab/tensors.py +++ b/python/dune/perftool/pdelab/tensors.py @@ -10,10 +10,14 @@ import pymbolic.primitives as prim import numpy as np -def define_list_tensor(name, expr, visitor): +def define_list_tensor(name, expr, visitor, stack=()): for i, child in enumerate(expr.ufl_operands): - instruction(assignee=prim.Subscript(prim.Variable(name), (i,)), - expression=visitor.call(child)) + from ufl.classes import ListTensor + if isinstance(child, ListTensor): + define_list_tensor(name, child, visitor, stack=stack + (i,)) + else: + instruction(assignee=prim.Subscript(prim.Variable(name), stack + (i,)), + expression=visitor.call(child)) @kernel_cached @@ -22,6 +26,7 @@ def pymbolic_list_tensor(expr, visitor): temporary_variable(name, shape=expr.ufl_shape, dtype=np.float64, + managed=True, ) define_list_tensor(name, expr, visitor) return prim.Variable(name)