Allow primal/dual error in Mult to make tests pass

Author: pbrubeck

firedrake/slate/slate.py

@@ -1267,7 +1267,7 @@ def __init__(self, A, B):
         fsA = A.arg_function_spaces[-1]
         fsB = B.arg_function_spaces[0]
 
-        assert fsA == fsB.dual(), (
+        assert space_equivalence(fsA, fsB), (
             "Cannot perform argument contraction over middle indices. "
             "They should be in dual function spaces."
         )
@@ -1317,7 +1317,7 @@ def __new__(cls, A, B, decomposition=None):
         fsA = A.arg_function_spaces[0]
         fsB = B.arg_function_spaces[0]
 
-        assert fsA == fsB, (
+        assert space_equivalence(fsA, fsB), (
             "Cannot perform argument contraction over middle indices. "
             "They must be in the same function space."
         )
@@ -1401,6 +1401,19 @@ def _output_string(self, prec=None):
         return "(%s).diag" % tensor
 
 
+def space_equivalence(A, B):
+    """Checks that two function spaces are equivalent.
+
+    :arg A: A function space.
+    :arg B: Another function space.
+
+    Returns `True` if they have matching meshes, elements, and rank. Otherwise,
+    `False` is returned.
+    """
+
+    return A.mesh() == B.mesh() and A.ufl_element() == B.ufl_element()
+
+
 def as_slate(F):
     """Convert an assembled or unassembled expression into a Slate Tensor.