LRE failure for chunking (#2608)

* change what's being asserted

* change assert, update docstring

* validate failure

* nate's mock test

* rename test; cleanup docstring

---------

Co-authored-by: nate stemen <nate@unitary.foundation>

Author: purva-thakre

mitiq/lre/tests/test_lre.py

@@ -10,6 +10,7 @@
 
 from mitiq import SUPPORTED_PROGRAM_TYPES, benchmarks
 from mitiq.lre import execute_with_lre, lre_decorator, mitigate_executor
+from mitiq.lre.multivariate_scaling.layerwise_folding import _get_chunks
 from mitiq.zne.scaling import fold_all, fold_global
 
 # default circuit for all unit tests
@@ -111,19 +112,6 @@ def execute(circuit, noise_level=0.025):
         )
 
 
-def test_lre_executor_with_chunking():
-    """Verify the executor works as expected for chunking a large circuit into
-    a smaller circuit."""
-    # define a larger circuit
-    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=12)[
-        0
-    ]
-    lre_exp_val = execute_with_lre(
-        test_cirq, execute, degree=2, fold_multiplier=2, num_chunks=14
-    )
-    assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
-
-
 @pytest.mark.parametrize("input_method", [(fold_global), (fold_all)])
 def test_lre_executor_with_different_folding_methods(input_method):
     """Verify the executor works as expected for using non-default unitary
@@ -136,3 +124,30 @@ def test_lre_executor_with_different_folding_methods(input_method):
         folding_method=input_method,
     )
     assert abs(lre_exp_val - ideal_val) <= abs(noisy_val - ideal_val)
+
+
+def test_lre_runs_correct_number_of_circuits_when_chunking():
+    """Verify execute_with_lre works as expected when chunking is used.
+    Note that this does not validate performance of chunking."""
+
+    mock_executor = Mock(side_effect=lambda _: random.random())
+
+    test_cirq = benchmarks.generate_rb_circuits(n_qubits=1, num_cliffords=12)[
+        0
+    ]
+
+    degree, fold_multiplier, num_chunks = 2, 2, 10
+
+    lre_exp_val_chunking = execute_with_lre(
+        test_cirq,
+        mock_executor,
+        degree=degree,
+        fold_multiplier=fold_multiplier,
+        num_chunks=num_chunks,
+    )
+
+    chunked_circ = _get_chunks(test_cirq, num_chunks=num_chunks)
+    assert isinstance(lre_exp_val_chunking, float)
+    assert mock_executor.call_count == math.comb(
+        degree + len(chunked_circ), degree
+    )