diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py
index 114651f793..f96b961218 100644
--- a/firedrake/formmanipulation.py
+++ b/firedrake/formmanipulation.py
@@ -144,11 +144,8 @@ def cofunction(self, o):
             # Not on a mixed space, just return ourselves.
             return o
 
-        try:
-            indices, = set(self.blocks.values())
-        except ValueError:
-            raise ValueError("Cofunction found on an off-diagonal block")
-
+        # We only need the test space for Cofunction		
+       indices = self.blocks[0]
         if len(indices) == 1:
             i = indices[0]
             W = V[i]
@@ -196,15 +193,15 @@ def split_form(form, diagonal=False):
     args = form.arguments()
     shape = tuple(len(a.function_space()) for a in args)
     forms = []
-    arity = len(shape)
+    rank = len(shape)
     if diagonal:
-        assert arity == 2
-        arity = 1
+        assert rank == 2
+        rank = 1
     for idx in numpy.ndindex(shape):
         if diagonal:
             i, j = idx
             if i != j:
                 continue
         f = splitter.split(form, idx)
-        forms.append(SplitForm(indices=idx[:arity], form=f))
+        forms.append(SplitForm(indices=idx[:rank], form=f))
     return tuple(forms)
diff --git a/firedrake/slate/slac/tsfc_driver.py b/firedrake/slate/slac/tsfc_driver.py
index 136b2a0084..0f5fbf96d3 100644
--- a/firedrake/slate/slac/tsfc_driver.py
+++ b/firedrake/slate/slac/tsfc_driver.py
@@ -50,7 +50,6 @@ def compile_terminal_form(tensor, prefix, *, tsfc_parameters=None):
     assert tensor.terminal, (
         "Only terminal tensors have forms associated with them!"
     )
-
     # Sets a default name for the subkernel prefix.
     mapper = RemoveRestrictions()
     integrals = map(partial(map_integrand_dags, mapper),