Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] `param_shift` with `broadcast=True` does not work with zero-length recipes
### Expected behavior

`param_shift` has feature parity between `broadcast=True` and `broadcast=False`

### Actual behavior

The (somewhat esoteric) example from the test suite below runs with `broadcast=False` but not with `broadcast=True`.

### Additional information

If we do not include the commented `assert` line, the gradient computation will raise an error, because the created tapes are not valid.

If we include the commented `assert` line in the code below, we see that too many tapes are being created when `broadcast=True`. This is because unnecessary tapes with `batch_size=0` are being created, which is the core bug making the tapes invalid.



### Source code

```shell
ops_with_custom_recipe = [1]
broadcast = True

dev = qml.device("default.qubit", wires=2)
x = [0.543, -0.654]

with qml.queuing.AnnotatedQueue() as q:
    qml.RX(x[0], wires=[0])
    qml.RX(x[1], wires=[0])
    qml.expval(qml.PauliZ(0))

tape = qml.tape.QuantumScript.from_queue(q)
gradient_recipes = tuple(
    [[-1e7, 1, 0], [1e7, 1, 0]] if i in ops_with_custom_recipe else None for i in range(2)
)
tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes, broadcast=broadcast)
num_ops_standard = (2 - len(ops_with_custom_recipe))
# assert len(tapes) == (1 if broadcast else 2) * num_ops_standard + (tape.num_params != num_ops_standard)
grad = fn(qml.execute(tapes, dev, None))
print(grad)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 19
     17 num_ops_standard = (2 - len(ops_with_custom_recipe))
     18 # assert len(tapes) == (1 if broadcast else 2) * num_ops_standard + (tape.num_params != num_ops_standard)
---> 19 grad = fn(qml.execute(tapes, dev, None))
     20 print(grad)

File ~/repos/pennylane/pennylane/workflow/execution.py:616, in execute(tapes, device, gradient_fn, interface, transform_program, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp)
    614 # Exiting early if we do not need to deal with an interface boundary
    615 if no_interface_boundary_required:
--> 616     results = inner_execute(tapes)
    617     return post_processing(results)
    619 _grad_on_execution = False

File ~/repos/pennylane/pennylane/workflow/execution.py:297, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    294 transformed_tapes, transform_post_processing = transform_program(tapes)
    296 if transformed_tapes:
--> 297     results = device_execution(transformed_tapes)
    298 else:
    299     results = ()

File ~/repos/pennylane/pennylane/devices/modifiers/simulator_tracking.py:30, in _track_execute.<locals>.execute(self, circuits, execution_config)
     28 @wraps(untracked_execute)
     29 def execute(self, circuits, execution_config=DefaultExecutionConfig):
---> 30     results = untracked_execute(self, circuits, execution_config)
     31     if isinstance(circuits, QuantumScript):
     32         batch = (circuits,)

File ~/repos/pennylane/pennylane/devices/modifiers/single_tape_support.py:32, in _make_execute.<locals>.execute(self, circuits, execution_config)
     30     is_single_circuit = True
     31     circuits = (circuits,)
---> 32 results = batch_execute(self, circuits, execution_config)
     33 return results[0] if is_single_circuit else results

File ~/repos/pennylane/pennylane/devices/default_qubit.py:594, in DefaultQubit.execute(self, circuits, execution_config)
    591 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    593 if max_workers is None:
--> 594     return tuple(
    595         _simulate_wrapper(
    596             c,
    597             {
    598                 "rng": self._rng,
    599                 "debugger": self._debugger,
    600                 "interface": interface,
    601                 "state_cache": self._state_cache,
    602                 "prng_key": _key,
    603             },
    604         )
    605         for c, _key in zip(circuits, prng_keys)
    606     )
    608 vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
    609 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/repos/pennylane/pennylane/devices/default_qubit.py:595, in <genexpr>(.0)
    591 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    593 if max_workers is None:
    594     return tuple(
--> 595         _simulate_wrapper(
    596             c,
    597             {
    598                 "rng": self._rng,
    599                 "debugger": self._debugger,
    600                 "interface": interface,
    601                 "state_cache": self._state_cache,
    602                 "prng_key": _key,
    603             },
    604         )
    605         for c, _key in zip(circuits, prng_keys)
    606     )
    608 vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
    609 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/repos/pennylane/pennylane/devices/default_qubit.py:842, in _simulate_wrapper(circuit, kwargs)
    841 def _simulate_wrapper(circuit, kwargs):
--> 842     return simulate(circuit, **kwargs)

File ~/repos/pennylane/pennylane/devices/qubit/simulate.py:292, in simulate(circuit, debugger, state_cache, **execution_kwargs)
    290 if state_cache is not None:
    291     state_cache[circuit.hash] = state
--> 292 return measure_final_state(circuit, state, is_state_batched, rng=rng, prng_key=meas_key)

File ~/repos/pennylane/pennylane/devices/qubit/simulate.py:213, in measure_final_state(circuit, state, is_state_batched, **execution_kwargs)
    210         raise TypeError("Native mid-circuit measurements are only supported with finite shots.")
    212     if len(circuit.measurements) == 1:
--> 213         return measure(circuit.measurements[0], state, is_state_batched=is_state_batched)
    215     return tuple(
    216         measure(mp, state, is_state_batched=is_state_batched) for mp in circuit.measurements
    217     )
    219 # finite-shot case

File ~/repos/pennylane/pennylane/devices/qubit/measure.py:233, in measure(measurementprocess, state, is_state_batched)
    220 def measure(
    221     measurementprocess: MeasurementProcess, state: TensorLike, is_state_batched: bool = False
    222 ) -> TensorLike:
    223     """Apply a measurement process to a state.
    224 
    225     Args:
   (...)
    231         Tensorlike: the result of the measurement
    232     """
--> 233     return get_measurement_function(measurementprocess, state)(
    234         measurementprocess, state, is_state_batched
    235     )

File ~/repos/pennylane/pennylane/devices/qubit/measure.py:72, in state_diagonalizing_gates(measurementprocess, state, is_state_batched)
     70 wires = Wires(range(total_indices))
     71 flattened_state = flatten_state(state, total_indices)
---> 72 return measurementprocess.process_state(flattened_state, wires)

File ~/repos/pennylane/pennylane/measurements/expval.py:142, in ExpectationMP.process_state(self, state, wire_order)
    140     return qml.math.squeeze(self.eigvals())
    141 with qml.queuing.QueuingManager.stop_recording():
--> 142     prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
    143 # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
    144 return self._calculate_expectation(prob)

File ~/repos/pennylane/pennylane/measurements/probs.py:238, in ProbabilityMP.process_state(self, state, wire_order)
    236 prob = qml.math.transpose(prob, desired_axes)
    237 # flatten and return probabilities
--> 238 return qml.math.reshape(prob, flat_shape)

File ~/venvs/dev/lib/python3.10/site-packages/autoray/autoray.py:80, in do(fn, like, *args, **kwargs)
     31 """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
     32 dispatch to retrieve ``fn`` based on whichever library defines the class of
     33 the ``args[0]``, or the ``like`` keyword argument if specified.
   (...)
     77     <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
     78 """
     79 backend = choose_backend(fn, *args, like=like, **kwargs)
---> 80 return get_lib_fn(backend, fn)(*args, **kwargs)

File ~/venvs/dev/lib/python3.10/site-packages/numpy/core/fromnumeric.py:285, in reshape(a, newshape, order)
    200 @array_function_dispatch(_reshape_dispatcher)
    201 def reshape(a, newshape, order='C'):
    202     """
    203     Gives a new shape to an array without changing its data.
    204 
   (...)
    283            [5, 6]])
    284     """
--> 285     return _wrapfunc(a, 'reshape', newshape, order=order)

File ~/venvs/dev/lib/python3.10/site-packages/numpy/core/fromnumeric.py:59, in _wrapfunc(obj, method, *args, **kwds)
     56     return _wrapit(obj, method, *args, **kwds)
     58 try:
---> 59     return bound(*args, **kwds)
     60 except TypeError:
     61     # A TypeError occurs if the object does have such a method in its
     62     # class, but its signature is not identical to that of NumPy's. This
   (...)
     66     # Call _wrapit from within the except clause to ensure a potential
     67     # exception has a traceback chain.
     68     return _wrapit(obj, method, *args, **kwds)

ValueError: cannot reshape array of size 0 into shape (0,newaxis)
```


### System information

```shell
pl dev
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Fix `param_shift(..., broadcast=True)` with zero-length gradient recipes
**Context:**

**Description of the Change:**
Introduces a special branch for the rare case of zero-length parameter-shift recipes.

**Benefits:**
Fixes #5622

**Possible Drawbacks:**

**Related GitHub Issues:**
Fixes #5622 

[sc-62416]

Patch:
diff --git a/doc/releases/changelog-0.36.0.md b/doc/releases/changelog-0.36.0.md
index 77826722f..c81e7c571 100644
--- a/doc/releases/changelog-0.36.0.md
+++ b/doc/releases/changelog-0.36.0.md
@@ -564,8 +564,9 @@
   [(#5610)](https://github.com/PennyLaneAI/pennylane/pull/5610)
 
 * Using shot vectors with `param_shift(... broadcast=True)` caused a bug. This combination is no longer supported
-  and will be added again in the next release.
+  and will be added again in the next release. Fixed a bug with custom gradient recipes that only consist of unshifted terms.
   [(#5612)](https://github.com/PennyLaneAI/pennylane/pull/5612)
+  [(#5623)](https://github.com/PennyLaneAI/pennylane/pull/5623)
 
 * Cast the keys of the `CountsMP` measurements returned `dynamic_one_shot` to the type produced by `MeasurementValue.concretize`.
   [(#5587)](https://github.com/PennyLaneAI/pennylane/pull/5587)
diff --git a/pennylane/gradients/general_shift_rules.py b/pennylane/gradients/general_shift_rules.py
index 07e2b9fce..cccb0f5a9 100644
--- a/pennylane/gradients/general_shift_rules.py
+++ b/pennylane/gradients/general_shift_rules.py
@@ -451,6 +451,9 @@ def generate_shifted_tapes(tape, index, shifts, multipliers=None, broadcast=Fals
             the ``batch_size`` of the returned tape matches the length of ``shifts``.
     """
 
+    if len(shifts) == 0:
+        return tuple()
+
     if multipliers is None:
         multipliers = np.ones_like(shifts)
 


Test patch:
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index 7a0efd917..0dd03d8fa 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -536,7 +536,8 @@ class TestParamShift:
     #     tapes, _ = qml.gradients.param_shift(circuit.tape, broadcast=broadcast)
     #     assert tapes == []
 
-    def test_with_gradient_recipes(self):
+    @pytest.mark.parametrize("broadcast", [True, False])
+    def test_with_gradient_recipes(self, broadcast):
         """Test that the function behaves as expected"""
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -549,18 +550,34 @@ class TestParamShift:
         tape = qml.tape.QuantumScript.from_queue(q)
         tape.trainable_params = {0, 2}
         gradient_recipes = ([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], [[1, 1, 1], [2, 2, 2], [3, 3, 3]])
-        tapes, _ = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        tapes, _ = param_shift(tape, gradient_recipes=gradient_recipes, broadcast=broadcast)
 
-        assert len(tapes) == 5
-        assert [t.batch_size for t in tapes] == [None] * 5
-        assert tapes[0].get_parameters(trainable_only=False) == [0.2 * 1.0 + 0.3, 2.0, 3.0, 4.0]
-        assert tapes[1].get_parameters(trainable_only=False) == [0.5 * 1.0 + 0.6, 2.0, 3.0, 4.0]
-        assert tapes[2].get_parameters(trainable_only=False) == [1.0, 2.0, 1 * 3.0 + 1, 4.0]
-        assert tapes[3].get_parameters(trainable_only=False) == [1.0, 2.0, 2 * 3.0 + 2, 4.0]
-        assert tapes[4].get_parameters(trainable_only=False) == [1.0, 2.0, 3 * 3.0 + 3, 4.0]
+        if broadcast:
+            assert len(tapes) == 2
+            assert [t.batch_size for t in tapes] == [2, 3]
+
+            shifted_batch = [0.2 * 1.0 + 0.3, 0.5 * 1.0 + 0.6]
+            tape_par = tapes[0].get_parameters(trainable_only=False)
+            assert np.allclose(tape_par[0], shifted_batch)
+            assert tape_par[1:] == [2.0, 3.0, 4.0]
+
+            shifted_batch = [1 * 3.0 + 1, 2 * 3.0 + 2, 3 * 3.0 + 3]
+            tape_par = tapes[1].get_parameters(trainable_only=False)
+            assert tape_par[:2] == [1.0, 2.0]
+            assert np.allclose(tape_par[2], shifted_batch)
+            assert tape_par[3:] == [4.0]
+        else:
+            assert len(tapes) == 5
+            assert [t.batch_size for t in tapes] == [None] * 5
+            assert tapes[0].get_parameters(trainable_only=False) == [0.2 * 1.0 + 0.3, 2.0, 3.0, 4.0]
+            assert tapes[1].get_parameters(trainable_only=False) == [0.5 * 1.0 + 0.6, 2.0, 3.0, 4.0]
+            assert tapes[2].get_parameters(trainable_only=False) == [1.0, 2.0, 1 * 3.0 + 1, 4.0]
+            assert tapes[3].get_parameters(trainable_only=False) == [1.0, 2.0, 2 * 3.0 + 2, 4.0]
+            assert tapes[4].get_parameters(trainable_only=False) == [1.0, 2.0, 3 * 3.0 + 3, 4.0]
 
+    @pytest.mark.parametrize("broadcast", [True, False])
     @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
-    def test_recycled_unshifted_tape(self, ops_with_custom_recipe):
+    def test_recycled_unshifted_tape(self, ops_with_custom_recipe, broadcast):
         """Test that if the gradient recipe has a zero-shift component, then
         the tape is executed only once using the current parameter
         values."""
@@ -577,23 +594,28 @@ class TestParamShift:
             [[-1e7, 1, 0], [1e7, 1, 1e-7]] if i in ops_with_custom_recipe else None
             for i in range(2)
         )
-        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        tapes, fn = param_shift(tape, gradient_recipes=gradient_recipes, broadcast=broadcast)
 
-        # two tapes per parameter that doesn't use a custom recipe,
+        # two (one with broadcast) tapes per parameter that doesn't use a custom recipe,
         # one tape per parameter that uses custom recipe,
         # plus one global call if at least one uses the custom recipe
-        num_ops_standard_recipe = tape.num_params - len(ops_with_custom_recipe)
-        assert len(tapes) == 2 * num_ops_standard_recipe + len(ops_with_custom_recipe) + 1
+        num_custom = len(ops_with_custom_recipe)
+        num_ops_standard_recipe = tape.num_params - num_custom
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == tapes_per_param * num_ops_standard_recipe + num_custom + 1
         # Test that executing the tapes and the postprocessing function works
         grad = fn(qml.execute(tapes, dev, None))
         assert qml.math.allclose(grad, -np.sin(x[0] + x[1]), atol=1e-5)
 
+    @pytest.mark.parametrize("broadcast", [False, True])
     @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
     @pytest.mark.parametrize("multi_measure", [False, True])
-    def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe, multi_measure):
+    def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe, multi_measure, broadcast):
         """Test that if the gradient recipe has a zero-shift component, then
         the tape is executed only once using the current parameter
         values."""
+        if multi_measure and broadcast:
+            pytest.skip("Multiple measurements are not supported with `broadcast=True` yet.")
         dev = qml.device("default.qubit", wires=2)
         x = [0.543, -0.654]
 
@@ -608,12 +630,13 @@ class TestParamShift:
         gradient_recipes = tuple(
             [[-1e7, 1, 0], [1e7, 1, 0]] if i in ops_with_custom_recipe else None for i in range(2)
         )
-        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        tapes, fn = param_shift(tape, gradient_recipes=gradient_recipes, broadcast=broadcast)
 
-        # two tapes per parameter that doesn't use a custom recipe,
+        # two (one with broadcast) tapes per parameter that doesn't use a custom recipe,
         # plus one global (unshifted) call if at least one uses the custom recipe
         num_ops_standard_recipe = tape.num_params - len(ops_with_custom_recipe)
-        assert len(tapes) == 2 * num_ops_standard_recipe + int(
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == tapes_per_param * num_ops_standard_recipe + int(
             tape.num_params != num_ops_standard_recipe
         )
         # Test that executing the tapes and the postprocessing function works
@@ -662,8 +685,9 @@ class TestParamShift:
         assert qml.math.allclose(grad[0], -np.sin(x[0] + x[1]), atol=1e-5)
         assert qml.math.allclose(grad[1], 0, atol=1e-5)
 
+    @pytest.mark.parametrize("broadcast", [True, False])
     @pytest.mark.parametrize("y_wire", [0, 1])
-    def test_f0_provided(self, y_wire):
+    def test_f0_provided(self, y_wire, broadcast):
         """Test that if the original tape output is provided, then
         the tape is not executed additionally at the current parameter
         values."""
@@ -677,7 +701,7 @@ class TestParamShift:
         tape = qml.tape.QuantumScript.from_queue(q)
         gradient_recipes = ([[-1e7, 1, 0], [1e7, 1, 1e7]],) * 2
         f0 = dev.execute(tape)
-        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes, f0=f0)
+        tapes, fn = param_shift(tape, gradient_recipes=gradient_recipes, f0=f0, broadcast=broadcast)
 
         # one tape per parameter that impacts the expval
         assert len(tapes) == 2 if y_wire == 0 else 1