Merge branch 'master' into sc-91994-ub24-up

Author: runora95

.github/CHANGELOG.md

@@ -43,6 +43,9 @@
   Updating ROCM from 5.7 to 6.2.4. 
   [(#1158)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1158)
 
+- Fix using Torch with `AmplitudeEmbedding` by applying `qml.broadcast_expand` before decomposition in the preprocessing stage. 
+  [(#1175)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1175)
+
 <h3>Internal changes ⚙️</h3>
 
 - Use local catalyst repository instead of fetching on Github CI.

pennylane_lightning/lightning_gpu/lightning_gpu.py

@@ -169,6 +169,7 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
 
     name = "adjoint + lightning.gpu"
     program.add_transform(no_sampling, name=name)
+    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(
         decompose,
         stopping_condition=_adjoint_ops,
@@ -180,7 +181,6 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
     program.add_transform(
         validate_measurements, analytic_measurements=adjoint_measurements, name=name
     )
-    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(validate_adjoint_trainable_params)
 
 

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -159,6 +159,7 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
 
     name = "adjoint + lightning.kokkos"
     program.add_transform(no_sampling, name=name)
+    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(
         decompose,
         stopping_condition=_adjoint_ops,
@@ -170,7 +171,6 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
     program.add_transform(
         validate_measurements, analytic_measurements=adjoint_measurements, name=name
     )
-    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(validate_adjoint_trainable_params)
 
 

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -163,6 +163,7 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
 
     name = "adjoint + lightning.qubit"
     program.add_transform(no_sampling, name=name)
+    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(
         decompose,
         stopping_condition=_adjoint_ops,
@@ -174,7 +175,6 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
     program.add_transform(
         validate_measurements, analytic_measurements=adjoint_measurements, name=name
     )
-    program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(validate_adjoint_trainable_params)
 
 

tests/lightning_base/test_device.py

@@ -102,7 +102,9 @@
 if not LightningDevice._CPP_BINARY_AVAILABLE:  # pylint: disable=protected-access
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
 
-fixture_params = itertools.product([3, None], [np.complex64, np.complex128])  # wires x c_dtype
+fixture_params = list(
+    itertools.product([3, None], [np.complex64, np.complex128])
+)  # wires x c_dtype
 
 
 @pytest.fixture(params=fixture_params)
@@ -191,6 +193,7 @@ def test_add_adjoint_transforms(self):
 
         name = f"adjoint + {device_name}"
         expected_program.add_transform(no_sampling, name=name)
+        expected_program.add_transform(qml.transforms.broadcast_expand)
         expected_program.add_transform(
             decompose,
             stopping_condition=_adjoint_ops,
@@ -204,7 +207,6 @@ def test_add_adjoint_transforms(self):
             analytic_measurements=adjoint_measurements,
             name=name,
         )
-        expected_program.add_transform(qml.transforms.broadcast_expand)
         expected_program.add_transform(validate_adjoint_trainable_params)
 
         actual_program = qml.transforms.core.TransformProgram()
@@ -699,6 +701,7 @@ def test_preprocess(self, adjoint):
         if adjoint:
             name = f"adjoint + {device_name}"
             expected_program.add_transform(no_sampling, name=name)
+            expected_program.add_transform(qml.transforms.broadcast_expand)
             expected_program.add_transform(
                 decompose,
                 stopping_condition=_adjoint_ops,
@@ -712,7 +715,6 @@ def test_preprocess(self, adjoint):
                 analytic_measurements=adjoint_measurements,
                 name=name,
             )
-            expected_program.add_transform(qml.transforms.broadcast_expand)
             expected_program.add_transform(validate_adjoint_trainable_params)
 
         gradient_method = "adjoint" if adjoint else None

tests/test_adjoint_jacobian.py

@@ -46,8 +46,10 @@
 
     kokkos_args += [InitializationSettings().set_num_threads(2)]
 
-fixture_params = itertools.product(
-    [np.complex64, np.complex128], kokkos_args, [None, 3]  # c_dtype x kokkos_args x wires
+fixture_params = list(
+    itertools.product(
+        [np.complex64, np.complex128], kokkos_args, [None, 3]  # c_dtype x kokkos_args x wires
+    )
 )
 
 
@@ -566,7 +568,12 @@ class TestAdjointJacobianQNode:
 
     @pytest.fixture(params=fixture_params)
     def dev(self, request):
-        return qml.device(device_name, wires=request.param[1], c_dtype=request.param[0])
+        params = request.param
+        if device_name == "lightning.kokkos":
+            return qml.device(
+                device_name, wires=params[2], c_dtype=params[0], kokkos_args=params[1]
+            )
+        return qml.device(device_name, wires=params[2], c_dtype=params[0])
 
     def test_qnode(self, mocker, dev):
         """Test that specifying diff_method allows the adjoint method to be selected"""
@@ -988,6 +995,31 @@ def f(params1, params2):
 
         assert np.allclose(grad_adjoint, grad_fd, atol=tol)
 
+    def test_torch_amplitude_embedding(self, dev):
+        """Test that the adjoint Jacobian works with AmplitudeEmbedding in a TorchLayer"""
+
+        torch = pytest.importorskip("torch")
+
+        # Define a simple circuit with AmplitudeEmbedding and BasicEntanglerLayers
+        n_qubits = 3
+        qml.AmplitudeEmbedding.ndim_params = (1,)
+
+        @qml.qnode(dev)
+        def qnode(inputs, weights):
+            qml.AmplitudeEmbedding(inputs, wires=range(n_qubits), normalize=True)
+            qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
+            return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]
+
+        n_layers = 1
+        weight_shapes = {"weights": (n_layers, n_qubits)}
+
+        qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)
+        clayer_1 = torch.nn.Linear(2, 8)
+        xs_test = torch.tensor([[-0.7380, 0.3596], [-0.1870, 0.9423]])
+
+        t = clayer_1(xs_test)
+        qlayer(t)
+
 
 def circuit_ansatz(params, wires):
     """Circuit ansatz containing all the parametrized gates"""