[BUG] Wire label issue in simple circuit with E2E sparse execution

[isaacdevlugt]

Expected behavior

The circuit below should run without error.

Actual behavior

Seems like a wire order / label issue.

Additional information

No response

Source code

@qml.qnode(qml.device('default.qubit', wires=2))
def circuit(state, mat):
    qml.StatePrep(state, wires=[0, 1])
    qml.QubitUnitary(mat, wires=[0, 1])
    return qml.expval(qml.Z(0))

state = scipy.sparse.csr_matrix([0, 0, 1, 0])
mat = scipy.sparse.eye(2**2)
circuit(state, mat)

Tracebacks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], line 9
      7 state = scipy.sparse.csr_matrix([0, 0, 1, 0])
      8 mat = scipy.sparse.eye(2**2)
----> 9 circuit(state, mat)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/workflow/qnode.py:881, in QNode.__call__(self, *args, **kwargs)
    878     from ._capture_qnode import capture_qnode  # pylint: disable=import-outside-toplevel
    880     return capture_qnode(self, *args, **kwargs)
--> 881 return self._impl_call(*args, **kwargs)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/workflow/qnode.py:854, in QNode._impl_call(self, *args, **kwargs)
    851 # Calculate the classical jacobians if necessary
    852 self._transform_program.set_classical_component(self, args, kwargs)
--> 854 res = qml.execute(
    855     (tape,),
    856     device=self.device,
    857     diff_method=self.diff_method,
    858     interface=self.interface,
    859     transform_program=self._transform_program,
    860     gradient_kwargs=self.gradient_kwargs,
    861     **self.execute_kwargs,
    862 )
    863 res = res[0]
    865 # convert result to the interface in case the qfunc has no parameters

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/workflow/execution.py:244, in execute(tapes, device, diff_method, interface, transform_program, grad_on_execution, cache, cachesize, max_diff, device_vjp, postselect_mode, mcm_method, gradient_kwargs, mcm_config, config, inner_transform)
    241 if transform_program.is_informative:
    242     return post_processing(tapes)
--> 244 results = run(tapes, device, config, inner_transform)
    245 return post_processing(results)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/workflow/run.py:286, in run(tapes, device, config, inner_transform_program)
    282 no_interface_boundary_required = (
    283     config.interface == Interface.NUMPY or config.gradient_method == "backprop"
    284 )
    285 if no_interface_boundary_required:
--> 286     results = inner_execute(tapes)
    287     return results
    289 # TODO: Prune once support for tf-autograph is dropped

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/workflow/run.py:251, in _make_inner_execute.<locals>.inner_execute(tapes)
    248 transformed_tapes, transform_post_processing = inner_transform(tapes)
    250 if transformed_tapes:
--> 251     results = device.execute(transformed_tapes, execution_config=execution_config)
    252 else:
    253     results = ()

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/modifiers/simulator_tracking.py:28, in _track_execute.<locals>.execute(self, circuits, execution_config)
     26 @wraps(untracked_execute)
     27 def execute(self, circuits, execution_config=DefaultExecutionConfig):
---> 28     results = untracked_execute(self, circuits, execution_config)
     29     if isinstance(circuits, QuantumScript):
     30         batch = (circuits,)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/modifiers/single_tape_support.py:30, in _make_execute.<locals>.execute(self, circuits, execution_config)
     28     is_single_circuit = True
     29     circuits = (circuits,)
---> 30 results = batch_execute(self, circuits, execution_config)
     31 return results[0] if is_single_circuit else results

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/default_qubit.py:717, in DefaultQubit.execute(self, circuits, execution_config)
    707     warnings.warn(
    708         (
    709             "Jitting executions with many circuits may have substantial classical overhead."
   (...)    712         UserWarning,
    713     )
    715 if max_workers is None:
--> 717     return tuple(
    718         _simulate_wrapper(
    719             c,
    720             {
    721                 "rng": self._rng,
    722                 "debugger": self._debugger,
    723                 "interface": interface,
    724                 "state_cache": self._state_cache,
    725                 "prng_key": _key,
    726                 "mcm_method": execution_config.mcm_config.mcm_method,
    727                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    728             },
    729         )
    730         for c, _key in zip(circuits, prng_keys)
    731     )
    733 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    734 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/default_qubit.py:718, in <genexpr>(.0)
    707     warnings.warn(
    708         (
    709             "Jitting executions with many circuits may have substantial classical overhead."
   (...)    712         UserWarning,
    713     )
    715 if max_workers is None:
    717     return tuple(
--> 718         _simulate_wrapper(
    719             c,
    720             {
    721                 "rng": self._rng,
    722                 "debugger": self._debugger,
    723                 "interface": interface,
    724                 "state_cache": self._state_cache,
    725                 "prng_key": _key,
    726                 "mcm_method": execution_config.mcm_config.mcm_method,
    727                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    728             },
    729         )
    730         for c, _key in zip(circuits, prng_keys)
    731     )
    733 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    734 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/default_qubit.py:1051, in _simulate_wrapper(circuit, kwargs)
   1050 def _simulate_wrapper(circuit, kwargs):
-> 1051     return simulate(circuit, **kwargs)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/qubit/simulate.py:357, in simulate(circuit, debugger, state_cache, **execution_kwargs)
    354     return tuple(results)
    356 ops_key, meas_key = jax_random_split(prng_key)
--> 357 state, is_state_batched = get_final_state(
    358     circuit, debugger=debugger, prng_key=ops_key, **execution_kwargs
    359 )
    360 if state_cache is not None:
    361     state_cache[circuit.hash] = state

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/qubit/simulate.py:190, in get_final_state(circuit, debugger, **execution_kwargs)
    188 if isinstance(op, MidMeasureMP):
    189     prng_key, key = jax_random_split(prng_key)
--> 190 state = apply_operation(
    191     op,
    192     state,
    193     is_state_batched=is_state_batched,
    194     debugger=debugger,
    195     prng_key=key,
    196     tape_shots=circuit.shots,
    197     **execution_kwargs,
    198 )
    199 # Handle postselection on mid-circuit measurements
    200 if isinstance(op, qml.Projector):

File /opt/homebrew/Cellar/python@3.11/3.11.11/Frameworks/Python.framework/Versions/3.11/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 ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/qubit/apply_operation.py:232, in apply_operation(op, state, is_state_batched, debugger, **_)
    166 @singledispatch
    167 def apply_operation(
    168     op: qml.operation.Operator,
   (...)    172     **_,
    173 ):
    174     """Apply and operator to a given state.
    175 
    176     Args:
   (...)    230 
    231     """
--> 232     return _apply_operation_default(op, state, is_state_batched, debugger)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/qubit/apply_operation.py:253, in _apply_operation_default(op, state, is_state_batched, debugger)
    250 """The default behaviour of apply_operation, accessed through the standard dispatch
    251 of apply_operation, as well as conditionally in other dispatches."""
    252 if op.has_sparse_matrix and not op.has_matrix:
--> 253     return apply_operation_csr_matrix(op, state, is_state_batched=is_state_batched)
    254 if (
    255     len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD
    256     and math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
    257 ) or (op.batch_size and is_state_batched):
    258     return apply_operation_einsum(op, state, is_state_batched=is_state_batched)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/devices/qubit/apply_operation.py:244, in apply_operation_csr_matrix(op, state, is_state_batched)
    242 num_wires = len(original_shape) - int(is_state_batched)
    243 full_state = math.reshape(state, [-1, 2**num_wires])  # expected: [batch_size, 2**num_wires]
--> 244 state_opT = full_state @ op.sparse_matrix(wire_order=range(num_wires)).T
    245 state_reshaped = math.reshape(state_opT, original_shape)
    246 return state_reshaped

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/operation.py:901, in Operator.sparse_matrix(self, wire_order, format)
    879 r"""Representation of the operator as a sparse matrix in the computational basis.
    880 
    881 If ``wire_order`` is provided, the numerical representation considers the position of the
   (...)    895 
    896 """
    897 canonical_sparse_matrix = self.compute_sparse_matrix(
    898     *self.parameters, format="csr", **self.hyperparameters
    899 )
--> 901 return expand_matrix(
    902     canonical_sparse_matrix, wires=self.wires, wire_order=wire_order
    903 ).asformat(format)

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/math/matrix_manipulation.py:145, in expand_matrix(mat, wires, wire_order, sparse_format)
    140     return qml.math.kron(mat1, mat2, like=interface)
    142 # get a subset of `wire_order` values that contain all wire labels inside `wires` argument
    143 # e.g. wire_order = [0, 1, 2, 3, 4]; wires = [3, 0, 2]
    144 # --> subset_wire_order = [0, 1, 2, 3]; expanded_wires = [3, 0, 2, 1]
--> 145 wire_indices = [wire_order.index(wire) for wire in wires]
    146 subset_wire_order = wire_order[min(wire_indices) : max(wire_indices) + 1]
    147 wire_difference = list(set(subset_wire_order) - set(wires))

File ~/.virtualenvs/pl-master/lib/python3.11/site-packages/pennylane/math/matrix_manipulation.py:145, in <listcomp>(.0)
    140     return qml.math.kron(mat1, mat2, like=interface)
    142 # get a subset of `wire_order` values that contain all wire labels inside `wires` argument
    143 # e.g. wire_order = [0, 1, 2, 3, 4]; wires = [3, 0, 2]
    144 # --> subset_wire_order = [0, 1, 2, 3]; expanded_wires = [3, 0, 2, 1]
--> 145 wire_indices = [wire_order.index(wire) for wire in wires]
    146 subset_wire_order = wire_order[min(wire_indices) : max(wire_indices) + 1]
    147 wire_difference = list(set(subset_wire_order) - set(wires))

ValueError: 1 is not in list

System information

Name: PennyLane
Version: 0.41.0.dev60
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: /Users/isaac/.virtualenvs/pl-master/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-15.3.2-arm64-arm-64bit
Python version:          3.11.11
Numpy version:           2.2.4
Scipy version:           1.15.2
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.40.0)
- default.clifford (PennyLane-0.41.0.dev60)
- default.gaussian (PennyLane-0.41.0.dev60)
- default.mixed (PennyLane-0.41.0.dev60)
- default.qubit (PennyLane-0.41.0.dev60)
- default.qutrit (PennyLane-0.41.0.dev60)
- default.qutrit.mixed (PennyLane-0.41.0.dev60)
- default.tensor (PennyLane-0.41.0.dev60)
- null.qubit (PennyLane-0.41.0.dev60)
- reference.qubit (PennyLane-0.41.0.dev60)

Existing GitHub issues

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