[BUG] shadow_expval doesn't work with parameter broadcasting

[CatalinaAlbornoz]

Expected behavior

You can do parameter broadcasting while having shadow_expval as the measurement.

Actual behavior

ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (100,3,3)->(100,3,3) (100,2,2)->(100,2,2)

Additional information

This also causes issues when using jax.vmap as highlighted in this Forum thread.

Source code

observables=[qml.PauliZ(0),qml.PauliZ(1),qml.PauliZ(2)]
dev = qml.device("default.qubit", wires=nq, shots=100)
@qml.qnode(dev) 
def qnode(x):
    qml.RX(x,0)
    return qml.shadow_expval(observables)

qnode([0.1, 0.2, 0.3])

Tracebacks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-14-76a6564bc821> in <cell line: 8>()
      6     return qml.shadow_expval(observables)
      7 
----> 8 qnode([0.1, 0.2, 0.3])

22 frames
/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in __call__(self, *args, **kwargs)
   1018         if qml.capture.enabled():
   1019             return qml.capture.qnode_call(self, *args, **kwargs)
-> 1020         return self._impl_call(*args, **kwargs)
   1021 
   1022 

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in _impl_call(self, *args, **kwargs)
   1006 
   1007         try:
-> 1008             res = self._execution_component(args, kwargs, override_shots=override_shots)
   1009         finally:
   1010             if old_interface == "auto":

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/qnode.py in _execution_component(self, args, kwargs, override_shots)
    955             )
    956             # pylint: disable=unexpected-keyword-arg
--> 957             res = qml.execute(
    958                 (self._tape,),
    959                 device=self.device,

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/execution.py in execute(tapes, device, gradient_fn, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp, mcm_config)
    769 
    770     if interface in jpc_interfaces:
--> 771         results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
    772     else:
    773         results = ml_boundary_execute(

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/interfaces/autograd.py in autograd_execute(tapes, execute_fn, jpc, device)
    145         [autograd.builtins.list(t.get_parameters()) for t in tapes]
    146     )
--> 147     return _execute(parameters, tuple(tapes), execute_fn, jpc)
    148 
    149 

/usr/local/lib/python3.10/dist-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/interfaces/autograd.py in _execute(parameters, tapes, execute_fn, jpc)
    166 
    167     """
--> 168     return execute_fn(tapes)
    169 
    170 

/usr/local/lib/python3.10/dist-packages/pennylane/workflow/execution.py in inner_execute(tapes, **_)
    210 
    211         if transformed_tapes:
--> 212             results = device.execute(transformed_tapes, execution_config=execution_config)
    213         else:
    214             results = ()

/usr/local/lib/python3.10/dist-packages/pennylane/devices/modifiers/simulator_tracking.py in 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,)

/usr/local/lib/python3.10/dist-packages/pennylane/devices/modifiers/single_tape_support.py in 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
     34 

/usr/local/lib/python3.10/dist-packages/pennylane/logging/decorators.py in wrapper_entry(*args, **kwargs)
     59                 **_debug_log_kwargs,
     60             )
---> 61         return func(*args, **kwargs)
     62 
     63     @wraps(func)

/usr/local/lib/python3.10/dist-packages/pennylane/devices/default_qubit.py in execute(self, circuits, execution_config)
    628 
    629         if max_workers is None:
--> 630             return tuple(
    631                 _simulate_wrapper(
    632                     c,

/usr/local/lib/python3.10/dist-packages/pennylane/devices/default_qubit.py in <genexpr>(.0)
    629         if max_workers is None:
    630             return tuple(
--> 631                 _simulate_wrapper(
    632                     c,
    633                     {

/usr/local/lib/python3.10/dist-packages/pennylane/devices/default_qubit.py in _simulate_wrapper(circuit, kwargs)
    894 
    895 def _simulate_wrapper(circuit, kwargs):
--> 896     return simulate(circuit, **kwargs)
    897 
    898 

/usr/local/lib/python3.10/dist-packages/pennylane/logging/decorators.py in wrapper_entry(*args, **kwargs)
     59                 **_debug_log_kwargs,
     60             )
---> 61         return func(*args, **kwargs)
     62 
     63     @wraps(func)

/usr/local/lib/python3.10/dist-packages/pennylane/devices/qubit/simulate.py in simulate(circuit, debugger, state_cache, **execution_kwargs)
    382     if state_cache is not None:
    383         state_cache[circuit.hash] = state
--> 384     return measure_final_state(
    385         circuit, state, is_state_batched, prng_key=meas_key, **execution_kwargs
    386     )

/usr/local/lib/python3.10/dist-packages/pennylane/logging/decorators.py in wrapper_entry(*args, **kwargs)
     59                 **_debug_log_kwargs,
     60             )
---> 61         return func(*args, **kwargs)
     62 
     63     @wraps(func)

/usr/local/lib/python3.10/dist-packages/pennylane/devices/qubit/simulate.py in measure_final_state(circuit, state, is_state_batched, **execution_kwargs)
    278     # finite-shot case
    279     rng = default_rng(rng)
--> 280     results = measure_with_samples(
    281         circuit.measurements,
    282         state,

/usr/local/lib/python3.10/dist-packages/pennylane/devices/qubit/sampling.py in measure_with_samples(measurements, state, shots, is_state_batched, rng, prng_key, mid_measurements)
    253         prng_key, key = jax_random_split(prng_key)
    254         all_res.extend(
--> 255             measure_fn(
    256                 group, state, shots, is_state_batched=is_state_batched, rng=rng, prng_key=key
    257             )

/usr/local/lib/python3.10/dist-packages/pennylane/devices/qubit/sampling.py in _measure_classical_shadow(mp, state, shots, is_state_batched, rng, prng_key)
    382         return [tuple(mp.process_state_with_shots(state, wires, s, rng=rng) for s in shots)]
    383 
--> 384     return [mp.process_state_with_shots(state, wires, shots.total_shots, rng=rng)]
    385 
    386 

/usr/local/lib/python3.10/dist-packages/pennylane/measurements/classical_shadow.py in process_state_with_shots(self, state, wire_order, shots, rng)
    542             float: The estimate of the expectation value.
    543         """
--> 544         bits, recipes = qml.classical_shadow(
    545             wires=self.wires, seed=self.seed
    546         ).process_state_with_shots(state, wire_order, shots, rng=rng)

/usr/local/lib/python3.10/dist-packages/pennylane/measurements/classical_shadow.py in process_state_with_shots(self, state, wire_order, shots, rng)
    410 
    411             # sample the observables on the first qubit
--> 412             probs = (np.einsum("abc,acb->a", first_qubit_state, obs[:, active_qubit]) + 1) / 2
    413             samples = bit_rng.random(size=probs.shape) > probs
    414             outcomes[:, active_qubit] = samples

/usr/local/lib/python3.10/dist-packages/numpy/core/einsumfunc.py in einsum(out, optimize, *operands, **kwargs)
   1369         if specified_out:
   1370             kwargs['out'] = out
-> 1371         return c_einsum(*operands, **kwargs)
   1372 
   1373     # Check the kwargs to avoid a more cryptic error later, without having to

ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (100,3,3)->(100,3,3) (100,2,2)->(100,2,2)

System information

Name: PennyLane
Version: 0.38.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.85+-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.26.4
Scipy version:           1.13.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.38.0)
- default.clifford (PennyLane-0.38.0)
- default.gaussian (PennyLane-0.38.0)
- default.mixed (PennyLane-0.38.0)
- default.qubit (PennyLane-0.38.0)
- default.qubit.autograd (PennyLane-0.38.0)
- default.qubit.jax (PennyLane-0.38.0)
- default.qubit.legacy (PennyLane-0.38.0)
- default.qubit.tf (PennyLane-0.38.0)
- default.qubit.torch (PennyLane-0.38.0)
- default.qutrit (PennyLane-0.38.0)
- default.qutrit.mixed (PennyLane-0.38.0)
- default.tensor (PennyLane-0.38.0)
- null.qubit (PennyLane-0.38.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.