diff --git a/python/dune/perftool/sumfact/accumulation.py b/python/dune/perftool/sumfact/accumulation.py
index db18095b209d5f4a3d42bf8cf38b92536bc4b819..2e3996f61d9ae7c0554454a12076ab7d766d3576 100644
--- a/python/dune/perftool/sumfact/accumulation.py
+++ b/python/dune/perftool/sumfact/accumulation.py
@@ -68,8 +68,42 @@ class AlreadyAssembledInput(SumfactKernelInputBase):
     def __eq__(self, other):
         return type(self) == type(other) and self.index == other.index
 
+    def __repr__(self):
+        return "AlreadyAssembledInput({})".format(self.index)
+
     def __hash__(self):
-        return 0
+        return hash(self.index)
+
+
+def _dimension_gradient_indices(argument):
+    """Return dimension and gradient index of test function argument
+
+    Dimension index corresponds to the dimensio for vector valued
+    functions (eg v in Stokes). The gradient index is the direction of
+    the derivative (eg in \Delat u in poisson or \grad u in Stoks).
+    """
+    grad_index = None
+    dimension_index = None
+
+    # If this is a gradient the last index is the grad_index
+    if argument.reference_grad:
+        grad_index = argument.expr.ufl_operands[1][-1]._value
+
+    # If this argument has indices there could be a dimension_index
+    if isinstance(argument.expr, uc.Indexed):
+        # More than two indices not supported
+        if len(argument.expr.ufl_operands[1]) > 2:
+            assert False
+        # For two indices the first is the dimension index
+        if len(argument.expr.ufl_operands[1]) == 2:
+            assert grad_index is not None
+            dimension_index = argument.expr.ufl_operands[1].indices()[0]._value
+        # For there is no gradient index we should have only one index, the dimension index
+        if not argument.reference_grad:
+            assert len(argument.expr.ufl_operands[1]) == 1
+            dimension_index = argument.expr.ufl_operands[1].indices()[0]._value
+
+    return dimension_index, grad_index
 
 
 @backend(interface="accum_insn", name="sumfact")
@@ -77,16 +111,16 @@ def generate_accumulation_instruction(visitor, accterm, measure, subdomain_id):
     # When doing sum factorization we want to split the test function
     assert(accterm.argument.expr is not None)
 
-    # For vector valued functions the first index is the dimension index
-    coeff_func_index = None
-    if isinstance(accterm.argument.expr, uc.Indexed):
-        coeff_func_index = accterm.argument.expr.ufl_operands[1].indices()[0]._value
+    # Get dimension and gradient index
+    dimension_index, grad_index = _dimension_gradient_indices(accterm.argument)
+    accum_index = (dimension_index,)
 
     # Do the tree traversal to get a pymbolic expression representing this expression
     pymbolic_expr = visitor(accterm.term)
     if pymbolic_expr == 0:
         return
 
+    # Number of basis functions
     dim = world_dimension()
     mod_arg_expr = accterm.argument.expr
     while (not isinstance(mod_arg_expr, uc.FunctionView)) and (not isinstance(mod_arg_expr, uc.Argument)):
@@ -95,20 +129,6 @@ def generate_accumulation_instruction(visitor, accterm, measure, subdomain_id):
     degree = mod_arg_expr.ufl_element()._degree
     basis_size = degree + 1
 
-    # The gradient index is the last
-    grad_index = None
-    if accterm.argument.reference_grad:
-        grad_index = accterm.argument.expr.ufl_operands[1][-1]._value
-
-    # Could be 0 so compare to None
-    accum_index = None
-    if coeff_func_index is not None:
-        accum_index = coeff_func_index
-    if accterm.argument.index and len(accterm.argument.index) > 1:
-        accum_index = accterm.argument.index[0]._value
-    if isinstance(accum_index, int):
-        accum_index = (accum_index,)
-
     jacobian_inames = tuple()
     if accterm.is_jacobian:
         jacobian_inames = visitor.inames
@@ -138,7 +158,7 @@ def generate_accumulation_instruction(visitor, accterm, measure, subdomain_id):
                            accumvar=accum,
                            within_inames=jacobian_inames,
                            input=AlreadyAssembledInput(index=accum_index),
-                           coeff_func_index=coeff_func_index,
+                           coeff_func_index=dimension_index,
                            )
 
         from dune.perftool.sumfact.vectorization import attach_vectorization_info