diff --git a/backend/src/plonky3/stark.rs b/backend/src/plonky3/stark.rs
index b6ec7422fc..7e0d0b679e 100644
--- a/backend/src/plonky3/stark.rs
+++ b/backend/src/plonky3/stark.rs
@@ -211,11 +211,17 @@ where
         let mut witness_by_machine = self
             .split
             .iter()
-            .map(|(machine, (pil, _))| {
-                (
-                    machine.clone(),
-                    machine_witness_columns(witness, pil, machine),
-                )
+            .filter_map(|(machine, (pil, _))| {
+                let witness_columns = machine_witness_columns(witness, pil, machine);
+                if witness_columns[0].1.is_empty() {
+                    // Empty machines can be removed entirely.
+                    None
+                } else {
+                    Some((
+                        machine.clone(),
+                        machine_witness_columns(witness, pil, machine),
+                    ))
+                }
             })
             .collect::<BTreeMap<_, _>>();
 
diff --git a/plonky3/src/prover.rs b/plonky3/src/prover.rs
index f0258a8d50..84558e1768 100644
--- a/plonky3/src/prover.rs
+++ b/plonky3/src/prover.rs
@@ -393,6 +393,9 @@ where
     }
 }
 
+/// Prove a program execution.
+/// Note that `witness_by_machine` might not have all the machines, empty ones are expected
+/// to be removed already.
 #[instrument(skip_all)]
 #[allow(clippy::multiple_bound_locations)] // cfg not supported in where clauses?
 pub fn prove<T: FieldElementMap>(
@@ -405,18 +408,12 @@ where
     ProverData<T>: Send,
     Commitment<T>: Send,
 {
-    let (tables, stage_0): (BTreeMap<_, _>, BTreeMap<_, _>) = program
-        .split
+    let (tables, stage_0): (BTreeMap<_, _>, BTreeMap<_, _>) = witness_by_machine
         .iter()
-        .filter_map(|(name, (_, constraint_system))| {
-            let columns = witness_by_machine.get(name).unwrap();
+        .map(|(name, columns)| {
+            let constraint_system = &program.split.get(name).unwrap().1;
             let degree = columns[0].1.len();
 
-            if degree == 0 {
-                // If a machine has no rows, remove it entirely.
-                return None;
-            }
-
             let table = Table {
                 air: PowdrTable::new(constraint_system),
                 degree,
@@ -433,7 +430,7 @@ where
                 );
             }
 
-            Some((
+            (
                 (name.clone(), table),
                 (
                     name.clone(),
@@ -455,7 +452,7 @@ where
                             .collect(),
                     },
                 ),
-            ))
+            )
         })
         .unzip();