Fixing kron compatibility numpy @ torch (#5540)

Fixes #5542

**Relevant Shortcut Stories:**
[sc-61616]

---------

Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>

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)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

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)
 
 

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):