diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index b8a5582c5b..8dfb774588 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -85,6 +85,8 @@ jobs:
             --no-binary h5py \
             --extra-index-url https://download.pytorch.org/whl/cpu \
             './firedrake-repo[ci]'
+          pip install -I "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@pbrubeck/remove-component-tensors"
+          pip install -I "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplify-indexed"
           firedrake-clean
           pip list
 
diff --git a/tests/firedrake/regression/test_vfs_component_bcs.py b/tests/firedrake/regression/test_vfs_component_bcs.py
index 1b636fefe4..9a5871041c 100644
--- a/tests/firedrake/regression/test_vfs_component_bcs.py
+++ b/tests/firedrake/regression/test_vfs_component_bcs.py
@@ -218,8 +218,8 @@ def test_component_full_bcs(V):
     A_cmp = assemble(a, bcs=bcs_cmp, mat_type="aij")
     A_mixed = assemble(a, bcs=bcs_mixed, mat_type="aij")
 
-    assert A_full.petscmat.equal(A_cmp.petscmat)
-    assert A_mixed.petscmat.equal(A_full.petscmat)
+    assert A_full.petscmat.equal(A_cmp.petscmat), str((A_full.petscmat[:, :]-A_cmp.petscmat[:, :]).tolist())
+    assert A_mixed.petscmat.equal(A_full.petscmat), str((A_full.petscmat[:, :]-A_mixed.petscmat[:, :]).tolist())
 
 
 def test_component_full_bcs_overlap(V):
diff --git a/tsfc/fem.py b/tsfc/fem.py
index 10c10dc97d..e20244a0a1 100644
--- a/tsfc/fem.py
+++ b/tsfc/fem.py
@@ -641,15 +641,13 @@ def fiat_to_ufl(fiat_dict, order):
 
 @translate.register(Argument)
 def translate_argument(terminal, mt, ctx):
-    argument_multiindex = ctx.argument_multiindices[terminal.number()]
-    sigma = tuple(gem.Index(extent=d) for d in mt.expr.ufl_shape)
     element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction)
 
     def callback(entity_id):
         finat_dict = ctx.basis_evaluation(element, mt, entity_id)
         # Filter out irrelevant derivatives
-        filtered_dict = {alpha: table
-                         for alpha, table in finat_dict.items()
+        filtered_dict = {alpha: finat_dict[alpha]
+                         for alpha in finat_dict
                          if sum(alpha) == mt.local_derivatives}
 
         # Change from FIAT to UFL arrangement
@@ -658,13 +656,16 @@ def callback(entity_id):
         # A numerical hack that FFC used to apply on FIAT tables still
         # lives on after ditching FFC and switching to FInAT.
         return ffc_rounding(square, ctx.epsilon)
+
     table = ctx.entity_selector(callback, mt.restriction)
     if ctx.use_canonical_quadrature_point_ordering:
         quad_multiindex = ctx.quadrature_rule.point_set.indices
         quad_multiindex_permuted = _make_quad_multiindex_permuted(mt, ctx)
         mapper = gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)
         table = mapper(table, tuple(zip(quad_multiindex, quad_multiindex_permuted)))
-    return gem.ComponentTensor(gem.Indexed(table, argument_multiindex + sigma), sigma)
+
+    argument_multiindex = ctx.argument_multiindices[terminal.number()]
+    return gem.partial_indexed(table, argument_multiindex)
 
 
 @translate.register(TSFCConstantMixin)