Repository: PennyLaneAI/pennylane

Issue:
[Bug] broadcasting in stochastic parameter shift gradients is broken
Seems to perform some batching checks that it shouldnt

```python
import pennylane as qml
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

H = jnp.polyval * qml.PauliY(0)

dev = qml.device("default.qubit", wires=range(1))
@qml.qnode(dev, interface="jax", diff_method=qml.gradients.stoch_pulse_grad, use_broadcasting=True)
def qnode(x):
    qml.evolve(H)((x[0],), t=1.)
    qml.evolve(H)((x[1],), t=1.)
    return qml.expval(qml.PauliZ(0))

x = jnp.ones((2, 2))
jax.grad(qnode)(x)
```

```pycon
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/Users/qottmann/Qottmann/Xanadu/Sandbox2/sandbox2.ipynb Cell 9 line 1
     13     return qml.expval(qml.PauliZ(0))
     15 x = jnp.ones((2, 2))
---> 16 jax.grad(qnode)(x)

    [... skipping hidden 10 frame]

File ~/Qottmann/Xanadu/pennylane/pennylane/qnode.py:1030, in QNode.__call__(self, *args, **kwargs)
   1025         full_transform_program._set_all_argnums(
   1026             self, args, kwargs, argnums
   1027         )  # pylint: disable=protected-access
   1029 # pylint: disable=unexpected-keyword-arg
-> 1030 res = qml.execute(
   1031     (self._tape,),
   1032     device=self.device,
   1033     gradient_fn=self.gradient_fn,
   1034     interface=self.interface,
   1035     transform_program=full_transform_program,
   1036     config=config,
   1037     gradient_kwargs=self.gradient_kwargs,
   1038     override_shots=override_shots,
   1039     **self.execute_kwargs,
   1040 )
   1042 res = res[0]
   1044 # convert result to the interface in case the qfunc has no parameters

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:786, 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)
    784     results = ml_boundary_execute(tapes, execute_fn, jpc)
    785 else:
--> 786     results = ml_boundary_execute(
    787         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
    788     )
    790 return post_processing(results)

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/jax.py:141, in execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff)
    132     return _execute_fwd(
    133         parameters,
    134         tapes,
   (...)
    137         _n=_n,
    138     )
    140 # PennyLane backward execution
--> 141 return _execute_bwd(
    142     parameters,
    143     tapes,
    144     device,
    145     execute_fn,
    146     gradient_fn,
    147     gradient_kwargs,
    148     _n=_n,
    149     max_diff=max_diff,
    150 )

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/jax.py:221, in _execute_bwd(params, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff)
    217         jvps = _compute_jvps(jacs, tangents[0], new_tapes)
    219     return res, jvps
--> 221 return execute_wrapper(params)

    [... skipping hidden 5 frame]

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/jax.py:195, in _execute_bwd.<locals>.execute_wrapper_jvp(primals, tangents)
    193 if at_max_diff:
    194     jvp_tapes, processing_fn = qml.gradients.batch_jvp(*_args, **_kwargs)
--> 195     jvps = processing_fn(execute_fn(jvp_tapes)[0])
    196 else:
    197     jvp_tapes, processing_fn = qml.gradients.batch_jvp(*_args, **_kwargs)

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:601, in execute.<locals>.inner_execute_with_empty_jac(tapes, **_)
    600 def inner_execute_with_empty_jac(tapes, **_):
--> 601     return (inner_execute(tapes), [])

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:252, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    250 if numpy_only:
    251     tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
--> 252 return cached_device_execution(tapes)

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:374, in cache_execute.<locals>.wrapper(tapes, **kwargs)
    369         return (res, []) if return_tuple else res
    371 else:
    372     # execute all unique tapes that do not exist in the cache
    373     # convert to list as new device interface returns a tuple
--> 374     res = list(fn(tuple(execution_tapes.values()), **kwargs))
    376 final_res = []
    378 for i, tape in enumerate(tapes):

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/default_qubit.py:478, in DefaultQubit.execute(self, circuits, execution_config)
    472 interface = (
    473     execution_config.interface
    474     if execution_config.gradient_method in {"backprop", None}
    475     else None
    476 )
    477 if max_workers is None:
--> 478     results = tuple(
    479         simulate(
    480             c,
    481             rng=self._rng,
    482             prng_key=self._prng_key,
    483             debugger=self._debugger,
    484             interface=interface,
    485         )
    486         for c in circuits
    487     )
    488 else:
    489     vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/default_qubit.py:479, in <genexpr>(.0)
    472 interface = (
    473     execution_config.interface
    474     if execution_config.gradient_method in {"backprop", None}
    475     else None
    476 )
    477 if max_workers is None:
    478     results = tuple(
--> 479         simulate(
    480             c,
    481             rng=self._rng,
    482             prng_key=self._prng_key,
    483             debugger=self._debugger,
    484             interface=interface,
    485         )
    486         for c in circuits
    487     )
    488 else:
    489     vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/qubit/simulate.py:228, in simulate(circuit, rng, prng_key, debugger, interface)
    197 def simulate(
    198     circuit: qml.tape.QuantumScript, rng=None, prng_key=None, debugger=None, interface=None
    199 ) -> Result:
    200     """Simulate a single quantum script.
    201 
    202     This is an internal function that will be called by the successor to ``default.qubit``.
   (...)
    226 
    227     """
--> 228     state, is_state_batched = get_final_state(circuit, debugger=debugger, interface=interface)
    229     return measure_final_state(circuit, state, is_state_batched, rng=rng, prng_key=prng_key)

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/qubit/simulate.py:123, in get_final_state(circuit, debugger, interface)
    121 is_state_batched = bool(prep and prep.batch_size is not None)
    122 for op in circuit.operations[bool(prep) :]:
--> 123     state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
    125     # Handle postselection on mid-circuit measurements
    126     if isinstance(op, qml.Projector):

File ~/anaconda3/envs/pennylane/lib/python3.11/functools.py:909, in singledispatch.<locals>.wrapper(*args, **kw)
    905 if not args:
    906     raise TypeError(f'{funcname} requires at least '
    907                     '1 positional argument')
--> 909 return dispatch(args[0].__class__)(*args, **kw)

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/qubit/apply_operation.py:338, in apply_parametrized_evolution(op, state, is_state_batched, debugger)
    335 """Apply ParametrizedEvolution by evolving the state rather than the operator matrix
    336 if we are operating on more than half of the subsystem"""
    337 if is_state_batched and op.batch_size is None:
--> 338     raise RuntimeError(
    339         "ParameterizedEvolution does not support standard broadcasting, but received a batched state"
    340     )
    342 # shape(state) is static (not a tracer), we can use an if statement
    343 num_wires = len(qml.math.shape(state))

RuntimeError: ParameterizedEvolution does not support standard broadcasting, but received a batched state
```

Latest master, latest jax


Discussion:
Note that this also seems broken on legacy:

```python3
import pennylane as qml
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

H = jnp.polyval * qml.PauliY(0)

dev = qml.device("default.qubit.jax", wires=range(1))
@qml.qnode(dev, interface="jax", diff_method=qml.gradients.stoch_pulse_grad, use_broadcasting=True)
def qnode(x):
    qml.evolve(H)((x[0],), t=1.)
    qml.evolve(H)((x[1],), t=1.)
    return qml.expval(qml.PauliZ(0))

x = jnp.ones((2, 2))
jax.grad(qnode)(x)
```

```con
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/Users/qottmann/Qottmann/Xanadu/Sandbox2/sandbox2.ipynb Cell 9 line 1
     13     return qml.expval(qml.PauliZ(0))
     15 x = jnp.ones((2, 2))
---> 16 jax.grad(qnode)(x)

    [... skipping hidden 10 frame]

File ~/Qottmann/Xanadu/pennylane/pennylane/qnode.py:1030, in QNode.__call__(self, *args, **kwargs)
   1025         full_transform_program._set_all_argnums(
   1026             self, args, kwargs, argnums
   1027         )  # pylint: disable=protected-access
   1029 # pylint: disable=unexpected-keyword-arg
-> 1030 res = qml.execute(
   1031     (self._tape,),
   1032     device=self.device,
   1033     gradient_fn=self.gradient_fn,
   1034     interface=self.interface,
   1035     transform_program=full_transform_program,
   1036     config=config,
   1037     gradient_kwargs=self.gradient_kwargs,
   1038     override_shots=override_shots,
   1039     **self.execute_kwargs,
   1040 )
...
   2571          "shape, got {} and {}.")
-> 2572   raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
   2574 return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)

TypeError: dot_general requires contracting dimensions to have the same shape, got (2,) and (6,).
```
Curiously, this seems to not be a problem when setting the device wires to `2` or more 🤔 
Hi @Qottmann! The device wires thing makes this interesting, but I'll share what I do know:

When creating a `ParametrizedEvolution` on [this line](https://github.com/PennyLaneAI/pennylane/blob/1f471160bf6ec57cbed949a6b8c339613869cc3f/pennylane/gradients/pulse_gradient.py#L128), the `before_t` value is a length-3 array (`Array([0. , 0.53697642, 1. ], dtype=float64)`), and `bcast`/`return_intermediate` is True. Perhaps it's set to True because the time value is unintentionally batched somewhere along the way? The `taus` passed into `_split_evol_ops` is a length-1 array on my local, so maybe that can be converted to a scalar somewhere and avoid it.

Anyway, assuming all this happens, we [set the batch size](https://github.com/PennyLaneAI/pennylane/blob/1f471160bf6ec57cbed949a6b8c339613869cc3f/pennylane/pulse/parametrized_evolution.py#L453) according to the time-batching, and as you saw in the error, we don't support batching.

WAIT! I think I remember intending on deleting that error for `ParametrizedEvolution` because we actually do support batching? The solution might simply be to delete the error being raised 😅 
ran the debugger, stopped it with the pre-processed circuit being executed and ran this:
```python
for op in circuit.operations:
    print((getattr(op, "t", None), op.batch_size, op))
```
the result:
```
(Array([0.        , 0.94974048, 1.        ], dtype=float64), 3, ParametrizedEvolution(array([1., 1.]), wires=[0]))
(None, None, PauliRot(1.5707963267948966, Y, wires=[0]))
(Array([0.        , 0.94974048, 1.        ], dtype=float64), 3, ParametrizedEvolution(array([1., 1.]), wires=[0]))
(array([0., 1.]), None, ParametrizedEvolution(array([1., 1.]), wires=[0]))
```
the last ParametrizedEvolution doesn't have a batch size, and its `t` value is indeed different. idk the details of why : )

Resolving patch:
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -334,23 +334,26 @@ def apply_parametrized_evolution(
 ):
     """Apply ParametrizedEvolution by evolving the state rather than the operator matrix
     if we are operating on more than half of the subsystem"""
-    if is_state_batched and op.batch_size is None:
-        raise RuntimeError(
-            "ParameterizedEvolution does not support standard broadcasting, but received a batched state"
-        )
 
     # shape(state) is static (not a tracer), we can use an if statement
-    num_wires = len(qml.math.shape(state))
+    num_wires = len(qml.math.shape(state)) - is_state_batched
     state = qml.math.cast(state, complex)
-    if 2 * len(op.wires) > num_wires and not op.hyperparameters["complementary"]:
-        # the subsystem operated is more than half of the system based on the state vector --> evolve state
-        return _evolve_state_vector_under_parametrized_evolution(op, state, num_wires)
-    # otherwise --> evolve matrix
-    return _apply_operation_default(op, state, is_state_batched, debugger)
+    if (
+        2 * len(op.wires) <= num_wires
+        or op.hyperparameters["complementary"]
+        or (is_state_batched and op.hyperparameters["return_intermediate"])
+    ):
+        # the subsystem operated on is half as big as the total system, or less
+        # or we want complementary time evolution
+        # or both the state and the operation have a batch dimension
+        # --> evolve matrix
+        return _apply_operation_default(op, state, is_state_batched, debugger)
+    # otherwise --> evolve state
+    return _evolve_state_vector_under_parametrized_evolution(op, state, num_wires, is_state_batched)
 
 
 def _evolve_state_vector_under_parametrized_evolution(
-    operation: qml.pulse.ParametrizedEvolution, state, num_wires
+    operation: qml.pulse.ParametrizedEvolution, state, num_wires, is_state_batched
 ):
     """Uses an odeint solver to compute the evolution of the input ``state`` under the given
     ``ParametrizedEvolution`` operation.
@@ -385,7 +388,13 @@ def _evolve_state_vector_under_parametrized_evolution(
             "You can update these values by calling the ParametrizedEvolution class: EV(params, t)."
         )
 
-    state = state.flatten()
+    if is_state_batched:
+        batch_dim = state.shape[0]
+        state = qml.math.moveaxis(state.reshape((batch_dim, 2**num_wires)), 1, 0)
+        out_shape = [2] * num_wires + [batch_dim]  # this shape is before moving the batch_dim back
+    else:
+        state = state.flatten()
+        out_shape = [2] * num_wires
 
     with jax.ensure_compile_time_eval():
         H_jax = ParametrizedHamiltonianPytree.from_hamiltonian(  # pragma: no cover
@@ -399,7 +408,9 @@ def fun(y, t):
         return (-1j * H_jax(operation.data, t=t)) @ y
 
     result = odeint(fun, state, operation.t, **operation.odeint_kwargs)
-    out_shape = [2] * num_wires
     if operation.hyperparameters["return_intermediate"]:
         return qml.math.reshape(result, [-1] + out_shape)
-    return qml.math.reshape(result[-1], out_shape)
+    result = qml.math.reshape(result[-1], out_shape)
+    if is_state_batched:
+        return qml.math.moveaxis(result, -1, 0)
+    return result


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -242,18 +242,11 @@ def test_globalphase(self, method, wire, ml_framework):
         assert qml.math.allclose(shift * initial_state, new_state_no_wire)
 
 
-# pylint:disable = unused-argument
 def time_independent_hamiltonian():
     """Create a time-independent Hamiltonian on two qubits."""
     ops = [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0), qml.PauliX(1)]
 
-    def f1(params, t):
-        return params  # constant
-
-    def f2(params, t):
-        return params  # constant
-
-    coeffs = [f1, f2, 4, 9]
+    coeffs = [qml.pulse.constant, qml.pulse.constant, 0.4, 0.9]
 
     return qml.pulse.ParametrizedHamiltonian(coeffs, ops)
 
@@ -275,10 +268,10 @@ def f2(params, t):
 
 
 @pytest.mark.jax
-class TestApplyParameterizedEvolution:
+class TestApplyParametrizedEvolution:
     @pytest.mark.parametrize("method", methods)
     def test_parameterized_evolution_time_independent(self, method):
-        """Test that applying a ParameterizedEvolution gives the expected state
+        """Test that applying a ParametrizedEvolution gives the expected state
         for a time-independent hamiltonian"""
 
         import jax.numpy as jnp
@@ -292,7 +285,7 @@ def test_parameterized_evolution_time_independent(self, method):
 
         H = time_independent_hamiltonian()
         params = jnp.array([1.0, 2.0])
-        t = 4
+        t = 0.4
 
         op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
 
@@ -306,7 +299,7 @@ def test_parameterized_evolution_time_independent(self, method):
 
     @pytest.mark.parametrize("method", methods)
     def test_parameterized_evolution_time_dependent(self, method):
-        """Test that applying a ParameterizedEvolution gives the expected state
+        """Test that applying a ParametrizedEvolution gives the expected state
         for a time dependent Hamiltonian"""
 
         import jax
@@ -321,7 +314,7 @@ def test_parameterized_evolution_time_dependent(self, method):
 
         H = time_dependent_hamiltonian()
         params = jnp.array([1.0, 2.0])
-        t = 4
+        t = 0.4
 
         op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
 
@@ -340,7 +333,7 @@ def generator(params):
         assert np.allclose(new_state, new_state_expected, atol=0.002)
 
     def test_large_state_small_matrix_evolves_matrix(self, mocker):
-        """Test that applying a ParameterizedEvolution operating on less
+        """Test that applying a ParametrizedEvolution operating on less
         than half of the wires in the state uses the default function to evolve
         the matrix"""
 
@@ -357,7 +350,7 @@ def test_large_state_small_matrix_evolves_matrix(self, mocker):
 
         H = time_independent_hamiltonian()
         params = jnp.array([1.0, 2.0])
-        t = 4
+        t = 0.4
 
         op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
 
@@ -374,7 +367,7 @@ def test_large_state_small_matrix_evolves_matrix(self, mocker):
         assert spy.call_count == 1
 
     def test_small_evolves_state(self, mocker):
-        """Test that applying a ParameterizedEvolution operating on less
+        """Test that applying a ParametrizedEvolution operating on less
         than half of the wires in the state uses the default function to evolve
         the matrix"""
 
@@ -433,7 +426,7 @@ def test_small_evolves_state(self, mocker):
 
         H = time_independent_hamiltonian()
         params = jnp.array([1.0, 2.0])
-        t = 4
+        t = 0.4
 
         op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
 
@@ -480,11 +473,16 @@ def test_parametrized_evolution_state_vector_return_intermediate(self, mocker):
         assert spy.call_count == 2
         assert qml.math.allclose(state_ev, state_rx, atol=1e-6)
 
-    def test_batched_state_raises_an_error(self):
-        """Test that if is_state_batche=True, an error is raised"""
+    @pytest.mark.parametrize("num_state_wires", [2, 4])
+    def test_with_batched_state(self, num_state_wires, mocker):
+        """Test that a ParametrizedEvolution is applied correctly to a batched state.
+        Note that the branching logic is different for batched input states, because
+        evolving the state vector does not support batching of the state. Instead,
+        the evolved matrix is used always."""
+        spy_einsum = mocker.spy(qml.math, "einsum")
         H = time_independent_hamiltonian()
         params = np.array([1.0, 2.0])
-        t = 4
+        t = 0.1
 
         op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
 
@@ -492,11 +490,24 @@ def test_batched_state_raises_an_error(self):
             [
                 [[0.81677345 + 0.0j, 0.0 + 0.0j], [0.0 - 0.57695852j, 0.0 + 0.0j]],
                 [[0.33894597 + 0.0j, 0.0 + 0.0j], [0.0 - 0.94080584j, 0.0 + 0.0j]],
+                [[0.33894597 + 0.0j, 0.0 + 0.0j], [0.0 - 0.94080584j, 0.0 + 0.0j]],
             ]
         )
+        if num_state_wires == 4:
+            zero_state_two_wires = np.eye(4)[0].reshape((2, 2))
+            initial_state = np.tensordot(initial_state, zero_state_two_wires, axes=0)
 
-        with pytest.raises(RuntimeError, match="does not support standard broadcasting"):
-            _ = apply_operation(op, initial_state, is_state_batched=True)
+        true_mat = qml.math.expm(-1j * qml.matrix(H(params, t=t)) * t)
+        U = qml.QubitUnitary(U=true_mat, wires=[0, 1])
+
+        new_state = apply_operation(op, initial_state, is_state_batched=True)
+        new_state_expected = apply_operation(U, initial_state, is_state_batched=True)
+        assert np.allclose(new_state, new_state_expected, atol=0.002)
+
+        if num_state_wires == 4:
+            assert spy_einsum.call_count == 2
+        else:
+            assert spy_einsum.call_count == 1
 
 
 @pytest.mark.parametrize("ml_framework", ml_frameworks_list)
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -134,7 +134,7 @@ def cost(x, cache):
 
 def atol_for_shots(shots):
     """Return higher tolerance if finite shots."""
-    return 2e-2 if shots else 1e-6
+    return 3e-2 if shots else 1e-6
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
diff --git a/tests/pulse/test_parametrized_evolution.py b/tests/pulse/test_parametrized_evolution.py
--- a/tests/pulse/test_parametrized_evolution.py
+++ b/tests/pulse/test_parametrized_evolution.py
@@ -776,7 +776,7 @@ def U(params):
 
 @pytest.mark.jax
 def test_map_wires():
-    """Test that map wires returns a new ParameterizedEvolution, with wires updated on
+    """Test that map wires returns a new ParametrizedEvolution, with wires updated on
     both the operator and the corresponding Hamiltonian"""
 
     def f1(p, t):