diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 780c64a2a25..b993ec77087 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -461,6 +461,9 @@ * Fixes a bug in `hamiltonian_expand` that produces incorrect output dimensions when shot vectors are combined with parameter broadcasting. [(#5494)](https://github.com/PennyLaneAI/pennylane/pull/5494) +* Fixes a bug in `qml.math.kron` that makes torch incompatible with numpy. + [(#5540)](https://github.com/PennyLaneAI/pennylane/pull/5540) +

Contributors ✍️

This release contains contributions from (in alphabetical order): diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 2c16ebe98e5..acf89bb4152 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -160,6 +160,14 @@ def kron(*args, like=None, **kwargs): """The kronecker/tensor product of args.""" if like == "scipy": return onp.kron(*args, **kwargs) # Dispatch scipy kron to numpy backed specifically. + + if like == "torch": + mats = [ + ar.numpy.asarray(arg, like="torch") if isinstance(arg, onp.ndarray) else arg + for arg in args + ] + return ar.numpy.kron(*mats) + return ar.numpy.kron(*args, like=like, **kwargs) diff --git a/tests/math/test_multi_dispatch.py b/tests/math/test_multi_dispatch.py index ac253f7c206..3c56c21f8a5 100644 --- a/tests/math/test_multi_dispatch.py +++ b/tests/math/test_multi_dispatch.py @@ -197,6 +197,17 @@ def test_dot_autograd(): assert fn.allclose(qml_grad(fn.dot)(x, y), x) +def test_kron(): + """Test the kronecker product function.""" + x = torch.tensor([[1, 2], [3, 4]]) + y = np.array([[0, 5], [6, 7]]) + + res = fn.kron(x, y) + expected = torch.tensor([[0, 5, 0, 10], [6, 7, 12, 14], [0, 15, 0, 20], [18, 21, 24, 28]]) + + assert fn.allclose(res, expected) + + class TestMatmul: @pytest.mark.torch def test_matmul_torch(self):