[BUG] Different ground states when batching with vmap or calling the function directly

erikrecio commented 3 months ago

Expected behavior

Both ground states gs[0] and g0 are the same.

Actual behavior

They are different. When calculated with a batch or directly calling the function.

Additional information

I'm not sure if it's a pennylane's problem or a jax one. I thouhgt to post it here first since the building of the matrix of the hamiltonian is all in pennylane.

Source code

import pennylane as qml
import jax
import numpy as np
import jax.numpy as jnp

def I(i):
    return qml.Identity(i)
def X(i):
    return qml.PauliX(i)
def Y(i):
    return qml.PauliY(i)
def Z(i):
    return qml.PauliZ(i)

nqubits = 4

def ground_state(j1, j2):

    hamiltonian = 0

    for i in range(nqubits):
        hamiltonian += Z(i)
        hamiltonian -= j1 * X(i) @ X((i+1)%nqubits)
        hamiltonian -= j2 * X((i-1)%nqubits) @ Z(i) @ X((i+1)%nqubits)

    ham_matrix = qml.matrix(hamiltonian)
    _, eigvecs = jnp.linalg.eigh(ham_matrix)

    return eigvecs[:,0]

j1 = 1
j2 = -3

gs = jax.vmap(ground_state, in_axes=[0,0])(np.array([j1]),np.array([j2]))
g0 = ground_state(j1, j2)

print((gs[0]==g0).all())

print(gs[0])
print(g0)

Tracebacks

No response

System information

Name: PennyLane
Version: 0.36.0.dev0
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: d:\documents\1. github\venv\lib\site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.10.11
Numpy version:           1.26.2
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.36.0.dev0)
- default.gaussian (PennyLane-0.36.0.dev0)
- default.mixed (PennyLane-0.36.0.dev0)
- default.qubit (PennyLane-0.36.0.dev0)
- default.qubit.autograd (PennyLane-0.36.0.dev0)
- default.qubit.jax (PennyLane-0.36.0.dev0)
- default.qubit.legacy (PennyLane-0.36.0.dev0)
- default.qubit.tf (PennyLane-0.36.0.dev0)
- default.qubit.torch (PennyLane-0.36.0.dev0)
- default.qutrit (PennyLane-0.36.0.dev0)
- null.qubit (PennyLane-0.36.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.34.0)

Existing GitHub issues

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

albi3ro commented 3 months ago

Thanks for opening this issue.

When trying it out locally, I get True for both master and v0.36.0, the latest release. But I do get False for version v0.35.0. We changed our system for representing arithmetic observables in v0.36.0.

Would you mind posting type(hamiltonian) and it's contents?

def ground_state(j1, j2):

    hamiltonian = 0

    for i in range(nqubits):
        hamiltonian += Z(i)
        hamiltonian -= j1 * X(i) @ X((i+1)%nqubits)
        hamiltonian -= j2 * X((i-1)%nqubits) @ Z(i) @ X((i+1)%nqubits)

    print(type(hamiltonian))
    print(hamiltonian)
    ham_matrix = qml.matrix(hamiltonian)
    _, eigvecs = jnp.linalg.eigh(ham_matrix)

    return eigvecs[:,0]

erikrecio commented 3 months ago

Seems like the new version fixes it! The types of hamiltonians before were different (from vmap to direct calling). Now they are both <class 'pennylane.ops.op_math.sum.Sum'>. The issues I was having seem to be fixed. Thanks for the quick answer!

albi3ro commented 3 months ago

Also, with the latest release, you can directly create qml.X ,qml.Y, qml.Z and qml.I.

albi3ro commented 3 months ago

Sorry. one last thing to mention, you may also want to do ham_matrix = qml.matrix(hamiltonian, wire_order=(0,1,2,3)). The new Sum operator gets it's wires from order of occurance.

erikrecio commented 3 months ago

Great, I'll take a look. Thanks!

PennyLaneAI / pennylane