Repository: PennyLaneAI/pennylane
Pull request #3852 (2023-03-01 20:39:18)

Add broadcasting support to `apply_operation` for new qubit SV device
**Description of the Change:**
* Added broadcasting support to the functions in `pennylane/devices/qubit/apply_operation.py`.
* `apply_operation_einsum` supports applying broadcasted operations to a state that is already broadcasted.
* Dispatched operator application functions have been updated as well.
* `apply_phase` has been removed.
* `apply_operation_tensordot` supports applying broadcasted operations to a non-broadcasted state, or non-broadcasted operations to broadcasted states, but not broadcasted operations to broadcasted state.
* `apply_operation` uses `apply_operation_einsum` if both the operation and state are broadcasted.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3852?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3852](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3852?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (bca3ca5) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/ce7704d4d2226ec6135ed8e2ac7739e3b38688ee?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ce7704d) will **decrease** coverage by `0.02%`.
> The diff coverage is `87.87%`.

```diff
@@            Coverage Diff             @@
##           master    #3852      +/-   ##
==========================================
- Coverage   99.71%   99.69%   -0.02%     
==========================================
  Files         338      338              
  Lines       29507    29516       +9     
==========================================
+ Hits        29422    29427       +5     
- Misses         85       89       +4     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3852?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/apply\_operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3852?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYXBwbHlfb3BlcmF0aW9uLnB5) | `94.66% <87.09%> (-5.34%)` | :arrow_down: |
| [pennylane/devices/qubit/simulate.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3852?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2ltdWxhdGUucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

mudit2812: [sc-33773]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index dad603f70..fbb8b17b1 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -29,6 +29,12 @@
 * `AdaptiveOptimizer` is updated to use non-default user-defined qnode arguments.
   [(#3765)](https://github.com/PennyLaneAI/pennylane/pull/3765)
 
+* The `apply_operation` function added to `devices/qubit` now supports broadcasting.
+  [(#3852)](https://github.com/PennyLaneAI/pennylane/pull/3852)
+
+* `qml.QubitStateVector.state_vector` now supports broadcasting.
+  [(#3852)](https://github.com/PennyLaneAI/pennylane/pull/3852)
+  
 * `pennylane.devices.qubit.preprocess` now allows circuits with non-commuting observables.
   [(#3857)](https://github.com/PennyLaneAI/pennylane/pull/3857)
 
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index cffc1d864..ec99d6d37 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Functions to apply an operation to a state vector."""
+# pylint: disable=unused-argument
+
 from functools import singledispatch
 from string import ascii_letters as alphabet
 
@@ -54,20 +56,21 @@ def _get_slice(index, axis, num_axes):
     return tuple(idx)
 
 
-def apply_operation_einsum(op: qml.operation.Operator, state):
+def apply_operation_einsum(op: qml.operation.Operator, state, is_state_batched: bool = False):
     """Apply ``Operator`` to ``state`` using ``einsum``. This is more efficent at lower qubit
     numbers.
 
     Args:
         op (Operator): Operator to apply to the quantum state
         state (array[complex]): Input quantum state
+        is_state_batched (bool): Boolean representing whether the state is batched or not
 
     Returns:
         array[complex]: output_state
     """
     mat = op.matrix()
 
-    total_indices = len(state.shape)
+    total_indices = len(state.shape) - is_state_batched
     num_indices = len(op.wires)
 
     state_indices = alphabet[:total_indices]
@@ -79,29 +82,44 @@ def apply_operation_einsum(op: qml.operation.Operator, state):
     for old, new in zip(affected_indices, new_indices):
         new_state_indices = new_state_indices.replace(old, new)
 
-    einsum_indices = f"{new_indices}{affected_indices},{state_indices}->{new_state_indices}"
+    einsum_indices = (
+        f"...{new_indices}{affected_indices},...{state_indices}->...{new_state_indices}"
+    )
 
-    reshaped_mat = math.reshape(mat, [2] * (num_indices * 2))
+    new_mat_shape = [2] * (num_indices * 2)
+    if op.batch_size is not None:
+        # Add broadcasting dimension to shape
+        new_mat_shape = [mat.shape[0]] + new_mat_shape
+    reshaped_mat = math.reshape(mat, new_mat_shape)
 
     return math.einsum(einsum_indices, reshaped_mat, state)
 
 
-def apply_operation_tensordot(op: qml.operation.Operator, state):
+def apply_operation_tensordot(op: qml.operation.Operator, state, is_state_batched: bool = False):
     """Apply ``Operator`` to ``state`` using ``math.tensordot``. This is more efficent at higher qubit
     numbers.
 
     Args:
         op (Operator): Operator to apply to the quantum state
         state (array[complex]): Input quantum state
+        is_state_batched (bool): Boolean representing whether the state is batched or not
 
     Returns:
         array[complex]: output_state
     """
     mat = op.matrix()
-    total_indices = len(state.shape)
+    total_indices = len(state.shape) - is_state_batched
     num_indices = len(op.wires)
-    reshaped_mat = math.reshape(mat, [2] * (num_indices * 2))
-    axes = (tuple(range(num_indices, 2 * num_indices)), tuple(op.wires))
+
+    new_mat_shape = [2] * (num_indices * 2)
+    if is_mat_batched := op.batch_size is not None:
+        # Add broadcasting dimension to shape
+        new_mat_shape = [mat.shape[0]] + new_mat_shape
+    reshaped_mat = math.reshape(mat, new_mat_shape)
+
+    mat_axes = list(range(-num_indices, 0))
+    state_axes = [i + is_state_batched for i in op.wires]
+    axes = (mat_axes, state_axes)
 
     tdot = math.tensordot(reshaped_mat, state, axes=axes)
 
@@ -111,18 +129,23 @@ def apply_operation_tensordot(op: qml.operation.Operator, state):
     # We'll need to invert this permutation to put the indices in the correct place
     unused_idxs = [i for i in range(total_indices) if i not in op.wires]
     perm = list(op.wires) + unused_idxs
-    inv_perm = math.argsort(perm)
+    if is_mat_batched:
+        perm = [0] + [i + 1 for i in perm]
+    if is_state_batched:
+        perm.insert(num_indices, -1)
 
+    inv_perm = math.argsort(perm)
     return math.transpose(tdot, inv_perm)
 
 
 @singledispatch
-def apply_operation(op: qml.operation.Operator, state):
+def apply_operation(op: qml.operation.Operator, state, is_state_batched: bool = False):
     """Apply and operator to a given state.
 
     Args:
         op (Operator): The operation to apply to ``state``
         state (tensor_like): The starting state.
+        is_state_batched (bool): Boolean representing whether the state is batched or not
 
     Returns:
         ndarray: output state
@@ -160,50 +183,43 @@ def apply_operation(op: qml.operation.Operator, state):
         [1., 0.]], requires_grad=True)
 
     """
-    if (len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD) and (
-        math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
-    ):
-        return apply_operation_einsum(op, state)
-    return apply_operation_tensordot(op, state)
+    if (
+        len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD
+        and math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
+    ) or (op.batch_size and is_state_batched):
+        return apply_operation_einsum(op, state, is_state_batched=is_state_batched)
+    return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
 
 
 @apply_operation.register
-def apply_paulix(op: qml.PauliX, state):
+def apply_paulix(op: qml.PauliX, state, is_state_batched: bool = False):
     """Apply :class:`pennylane.PauliX` operator to the quantum state"""
-    return math.roll(state, 1, op.wires[0])
+    axis = op.wires[0] + is_state_batched
+    return math.roll(state, 1, axis)
 
 
 @apply_operation.register
-def apply_pauliz(op: qml.PauliZ, state):
+def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False):
     """Apply pauliz to state."""
-    n_wires = math.ndim(state)
-    sl_0 = _get_slice(0, op.wires[0], n_wires)
-    sl_1 = _get_slice(1, op.wires[0], n_wires)
-
-    state1 = math.multiply(state[sl_1], -1)
-    return math.stack([state[sl_0], state1], axis=op.wires[0])
-
-
-@apply_operation.register
-def apply_phase(op: qml.PhaseShift, state):
-    """Apply PhaseShift operator to state."""
-    shift = math.exp(math.multiply(1j, op.data[0]))
+    axis = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
 
-    n_wires = math.ndim(state)
-    sl_0 = _get_slice(0, op.wires[0], n_wires)
-    sl_1 = _get_slice(1, op.wires[0], n_wires)
+    sl_0 = _get_slice(0, axis, n_dim)
+    sl_1 = _get_slice(1, axis, n_dim)
 
-    state1 = math.multiply(shift, state[sl_1])
-    return math.stack([state[sl_0], state1], axis=op.wires[0])
+    state1 = math.multiply(state[sl_1], -1)
+    return math.stack([state[sl_0], state1], axis=axis)
 
 
 @apply_operation.register
-def apply_cnot(op: qml.CNOT, state):
+def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False):
     """Apply cnot gate to state."""
-    target_axes = (op.wires[1] - 1) if op.wires[1] > op.wires[0] else (op.wires[1])
+    target_axes = (op.wires[1] - 1 if op.wires[1] > op.wires[0] else op.wires[1]) + is_state_batched
+    control_axes = op.wires[0] + is_state_batched
+    n_dim = math.ndim(state)
+
+    sl_0 = _get_slice(0, control_axes, n_dim)
+    sl_1 = _get_slice(1, control_axes, n_dim)
 
-    n_wires = math.ndim(state)
-    sl_0 = _get_slice(0, op.wires[0], n_wires)
-    sl_1 = _get_slice(1, op.wires[0], n_wires)
     state_x = math.roll(state[sl_1], 1, target_axes)
-    return math.stack([state[sl_0], state_x], axis=op.wires[0])
+    return math.stack([state[sl_0], state_x], axis=control_axes)
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 41c6feae8..46cdf2bbc 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -189,7 +189,8 @@ class QubitStateVector(StatePrep):
 
     def state_vector(self, wire_order=None):
         num_op_wires = len(self.wires)
-        op_vector = math.reshape(self.parameters[0], (2,) * num_op_wires)
+        op_vector_shape = (-1,) + (2,) * num_op_wires if self.batch_size else (2,) * num_op_wires
+        op_vector = math.reshape(self.parameters[0], op_vector_shape)
 
         if wire_order is None or Wires(wire_order) == self.wires:
             return op_vector
@@ -199,14 +200,26 @@ class QubitStateVector(StatePrep):
             raise WireError("Custom wire_order must contain all QubitStateVector wires")
 
         num_total_wires = len(wire_order)
-        indices = tuple([slice(None)] * num_op_wires + [0] * (num_total_wires - num_op_wires))
-        ket = np.zeros((2,) * num_total_wires, dtype=np.complex128)
+        indices = tuple(
+            [Ellipsis] + [slice(None)] * num_op_wires + [0] * (num_total_wires - num_op_wires)
+        )
+        ket_shape = [2] * num_total_wires
+        if self.batch_size:
+            # Add broadcasted dimension to the shape of the state vector
+            ket_shape = [self.batch_size] + ket_shape
+
+        ket = np.zeros(ket_shape, dtype=np.complex128)
         ket[indices] = op_vector
 
         # unless wire_order is [*self.wires, *rest_of_wire_order], need to rearrange
         if self.wires != wire_order[:num_op_wires]:
             current_order = self.wires + list(Wires.unique_wires([wire_order, self.wires]))
             desired_order = [current_order.index(w) for w in wire_order]
+            if self.batch_size:
+                # If the operation is broadcasted, the desired order must include the batch dimension
+                # as the first dimension.
+                desired_order = [0] + [d + 1 for d in desired_order]
+
             ket = ket.transpose(desired_order)
 
         return math.convert_like(ket, op_vector)


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 034b7db39..c3ecdf938 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -17,6 +17,7 @@ Tests the apply_operation functions from devices/qubit
 import pytest
 
 import numpy as np
+from scipy.stats import unitary_group
 import pennylane as qml
 
 
@@ -35,6 +36,9 @@ ml_frameworks_list = [
 ]
 
 
+methods = [apply_operation_einsum, apply_operation_tensordot, apply_operation]
+
+
 def test_custom_operator_with_matrix():
     """Test that apply_operation works with any operation that defines a matrix."""
 
@@ -59,9 +63,7 @@ def test_custom_operator_with_matrix():
 
 
 @pytest.mark.parametrize("ml_framework", ml_frameworks_list)
-@pytest.mark.parametrize(
-    "method", (apply_operation_einsum, apply_operation_tensordot, apply_operation)
-)
+@pytest.mark.parametrize("method", methods)
 @pytest.mark.parametrize("wire", (0, 1))
 class TestTwoQubitStateSpecialCases:
     """Test the special cases on a two qubit state.  Also tests the special cases for einsum and tensor application methods
@@ -204,9 +206,7 @@ class TestTwoQubitStateSpecialCases:
         assert qml.math.allclose(initial1[0], new1[1])
 
 
-@pytest.mark.parametrize(
-    "method", (apply_operation_einsum, apply_operation_tensordot, apply_operation)
-)
+@pytest.mark.parametrize("method", methods)
 class TestRXCalcGrad:
     """Tests the application and differentiation of an RX gate in the different interfaces."""
 
@@ -324,9 +324,84 @@ class TestRXCalcGrad:
         self.compare_expected_result(phi, state, new_state, phi_grad)
 
 
-@pytest.mark.parametrize(
-    "method", (apply_operation_einsum, apply_operation_tensordot, apply_operation)
-)
+@pytest.mark.parametrize("ml_framework", ml_frameworks_list)
+@pytest.mark.parametrize("method", methods)
+class TestBroadcasting:  # pylint: disable=too-few-public-methods
+    """Tests that broadcasted operations are applied correctly."""
+
+    broadcasted_ops = [
+        qml.RX(np.array([np.pi, np.pi / 2, np.pi / 4]), wires=2),
+        qml.PhaseShift(np.array([np.pi, np.pi / 2, np.pi / 4]), wires=2),
+        qml.IsingXX(np.array([np.pi, np.pi / 2, np.pi / 4]), wires=[1, 2]),
+        qml.QubitUnitary(
+            np.array([unitary_group.rvs(8), unitary_group.rvs(8), unitary_group.rvs(8)]),
+            wires=[0, 1, 2],
+        ),
+    ]
+
+    unbroadcasted_ops = [
+        qml.PauliX(2),
+        qml.PauliZ(2),
+        qml.CNOT([1, 2]),
+        qml.RX(np.pi, wires=2),
+        qml.PhaseShift(np.pi / 2, wires=2),
+        qml.IsingXX(np.pi / 2, wires=[1, 2]),
+        qml.QubitUnitary(unitary_group.rvs(8), wires=[0, 1, 2]),
+    ]
+
+    @pytest.mark.parametrize("op", broadcasted_ops)
+    def test_broadcasted_op(self, op, method, ml_framework):
+        """Tests that batched operations are applied correctly to an unbatched state."""
+        state = np.ones((2, 2, 2)) / np.sqrt(8)
+
+        res = method(op, qml.math.asarray(state, like=ml_framework))
+        missing_wires = 3 - len(op.wires)
+        mat = op.matrix()
+        expanded_mat = [
+            np.kron(np.eye(2**missing_wires), mat[i]) if missing_wires else mat[i]
+            for i in range(3)
+        ]
+        expected = [(expanded_mat[i] @ state.flatten()).reshape((2, 2, 2)) for i in range(3)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize("op", unbroadcasted_ops)
+    def test_broadcasted_state(self, op, method, ml_framework):
+        """Tests that unbatched operations are applied correctly to a batched state."""
+        state = np.ones((3, 2, 2, 2)) / np.sqrt(8)
+
+        res = method(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
+        missing_wires = 3 - len(op.wires)
+        mat = op.matrix()
+        expanded_mat = np.kron(np.eye(2**missing_wires), mat) if missing_wires else mat
+        expected = [(expanded_mat @ state[i].flatten()).reshape((2, 2, 2)) for i in range(3)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+    @pytest.mark.parametrize("op", broadcasted_ops)
+    def test_broadcasted_op_broadcasted_state(self, op, method, ml_framework):
+        """Tests that batched operations are applied correctly to a batched state."""
+        if method is apply_operation_tensordot:
+            pytest.skip("Tensordot doesn't support batched operator and batched state.")
+
+        state = np.ones((3, 2, 2, 2)) / np.sqrt(8)
+
+        res = method(op, qml.math.asarray(state, like=ml_framework), is_state_batched=True)
+        missing_wires = 3 - len(op.wires)
+        mat = op.matrix()
+        expanded_mat = [
+            np.kron(np.eye(2**missing_wires), mat[i]) if missing_wires else mat[i]
+            for i in range(3)
+        ]
+        expected = [(expanded_mat[i] @ state[i].flatten()).reshape((2, 2, 2)) for i in range(3)]
+
+        assert qml.math.get_interface(res) == ml_framework
+        assert qml.math.allclose(res, expected)
+
+
+@pytest.mark.parametrize("method", methods)
 class TestLargerOperations:
     """Tests matrix applications on states and operations with larger numbers of wires."""
 
diff --git a/tests/devices/qubit/test_initialize_state.py b/tests/devices/qubit/test_initialize_state.py
index de568a1a6..5aa1ad5aa 100644
--- a/tests/devices/qubit/test_initialize_state.py
+++ b/tests/devices/qubit/test_initialize_state.py
@@ -58,6 +58,33 @@ class TestInitializeState:
         state[0, 1, 0] = 0  # set to zero to make test below simple
         assert qml.math.allequal(state, np.zeros((2, 2, 2)))
 
+    def test_create_initial_state_with_QubitStateVector(self):
+        """Tests that create_initial_state works with the QubitStateVector operator."""
+        prep_op = qml.QubitStateVector(
+            np.array([0, 1, 0, 0, 0, 0, 0, 1]) / np.sqrt(2), wires=[0, 1, 2]
+        )
+        state = create_initial_state([0, 1, 2], prep_operation=prep_op)
+        assert state[0, 0, 1] == 1 / np.sqrt(2)
+        assert state[1, 1, 1] == 1 / np.sqrt(2)
+        state[0, 0, 1] = 0
+        state[1, 1, 1] = 0  # set to zero to make test below simple
+        assert qml.math.allequal(state, np.zeros((2, 2, 2)))
+
+    def test_create_initial_state_with_QubitStateVector_broadcasted(self):
+        """Tests that create_initial_state works with a broadcasted QubitStateVector
+        operator."""
+        prep_op = qml.QubitStateVector(
+            np.array([[0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]), wires=[0, 1]
+        )
+        state = create_initial_state([0, 1], prep_operation=prep_op)
+        assert state[0, 0, 1] == 1
+        assert state[1, 1, 1] == 1
+        assert state[2, 1, 0] == 1
+        state[0, 0, 1] = 0
+        state[1, 1, 1] = 0
+        state[2, 1, 0] = 0  # set to zero to make test below simple
+        assert qml.math.allequal(state, np.zeros((3, 2, 2)))
+
     def test_create_initialize_state_prefers_op_interface_over_like(self):
         """Tests that the like argument is ignored when a prep-op is provided."""
         prep_op = self.DefaultPrep([0, 0, 0, 1], wires=[0, 1])
diff --git a/tests/ops/qubit/test_state_prep.py b/tests/ops/qubit/test_state_prep.py
index 10655a0b8..551796b7a 100644
--- a/tests/ops/qubit/test_state_prep.py
+++ b/tests/ops/qubit/test_state_prep.py
@@ -119,6 +119,28 @@ class TestStateVector:
         ket[one_position] = 0  # everything else should be zero, as we assert below
         assert np.allclose(np.zeros((2,) * num_wires), ket)
 
+    @pytest.mark.parametrize(
+        "num_wires,wire_order,one_positions",
+        [
+            (2, None, [(0, 1, 0), (1, 0, 1)]),
+            (2, [1, 2], [(0, 1, 0), (1, 0, 1)]),
+            (3, [0, 1, 2], [(0, 0, 1, 0), (1, 0, 0, 1)]),
+            (3, ["a", 1, 2], [(0, 0, 1, 0), (1, 0, 0, 1)]),
+            (3, [1, 2, 0], [(0, 1, 0, 0), (1, 0, 1, 0)]),
+            (3, [1, 2, "a"], [(0, 1, 0, 0), (1, 0, 1, 0)]),
+            (3, [2, 1, 0], [(0, 0, 1, 0), (1, 1, 0, 0)]),
+            (4, [3, 2, 0, 1], [(0, 0, 0, 0, 1), (1, 0, 1, 0, 0)]),
+        ],
+    )
+    def test_QubitStateVector_state_vector_broadcasted(self, num_wires, wire_order, one_positions):
+        """Tests that QubitStateVector state_vector returns kets with broadcasting as expected."""
+        qsv_op = qml.QubitStateVector([[0, 0, 1, 0], [0, 1, 0, 0]], wires=[1, 2])  # |10>, |01>
+        ket = qsv_op.state_vector(wire_order=wire_order)
+        assert ket[one_positions[0]] == 1 == ket[one_positions[1]]
+        ket[one_positions[0]] = ket[one_positions[1]] = 0
+        # everything else should be zero, as we assert below
+        assert np.allclose(np.zeros((2,) * (num_wires + 1)), ket)
+
     def test_QubitStateVector_reordering(self):
         """Tests that wires get re-ordered as expected."""
         qsv_op = qml.QubitStateVector(np.array([1, -1, 1j, -1j]) / 2, wires=[0, 1])
@@ -127,6 +149,17 @@ class TestStateVector:
         expected[0, :, 0, :] = np.array([[1, 1j], [-1, -1j]]) / 2
         assert np.array_equal(ket, expected)
 
+    def test_QubitStateVector_reordering_broadcasted(self):
+        """Tests that wires get re-ordered as expected with broadcasting."""
+        qsv_op = qml.QubitStateVector(
+            np.array([[1, -1, 1j, -1j], [1, -1j, -1, 1j]]) / 2, wires=[0, 1]
+        )
+        ket = qsv_op.state_vector(wire_order=[2, 1, 3, 0])
+        expected = np.zeros((2,) * 5, dtype=np.complex128)
+        expected[0, 0, :, 0, :] = np.array([[1, 1j], [-1, -1j]]) / 2
+        expected[1, 0, :, 0, :] = np.array([[1, -1], [-1j, 1j]]) / 2
+        assert np.array_equal(ket, expected)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
     def test_QubitStateVector_state_vector_preserves_parameter_type(self, interface):
@@ -135,6 +168,16 @@ class TestStateVector:
         assert qml.math.get_interface(qsv_op.state_vector()) == interface
         assert qml.math.get_interface(qsv_op.state_vector(wire_order=[0, 1, 2])) == interface
 
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_QubitStateVector_state_vector_preserves_parameter_type_broadcasted(self, interface):
+        """Tests that given an array of some type, the resulting state vector is also that type."""
+        qsv_op = qml.QubitStateVector(
+            qml.math.array([[0, 0, 0, 1], [1, 0, 0, 0]], like=interface), wires=[1, 2]
+        )
+        assert qml.math.get_interface(qsv_op.state_vector()) == interface
+        assert qml.math.get_interface(qsv_op.state_vector(wire_order=[0, 1, 2])) == interface
+
     def test_QubitStateVector_state_vector_bad_wire_order(self):
         """Tests that the provided wire_order must contain the wires in the operation."""
         qsv_op = qml.QubitStateVector([0, 0, 0, 1], wires=[0, 1])