Catalyst treats a density matrix measurement as a state measurement

[albi3ro]

Issue description

Description of the issue - include code snippets and screenshots here if relevant. You may use the following template below

• Expected behavior: (What you expect to happen)

I would the expect the same behavior as in a non-qjit execution. Requesting qml.density_matrix(wires) should return a density matrix.

Without the qjit, I get:

array([[0.46939564+0.j        , 0.        +0.11985638j,
        0.46939564+0.j        , 0.        +0.11985638j],
       [0.        -0.11985638j, 0.03060436+0.j        ,
        0.        -0.11985638j, 0.03060436+0.j        ],
       [0.46939564+0.j        , 0.        +0.11985638j,
        0.46939564+0.j        , 0.        +0.11985638j],
       [0.        -0.11985638j, 0.03060436+0.j        ,
        0.        -0.11985638j, 0.03060436+0.j        ]])

• Actual behavior: (What actually happens)

We just get the state.

• Reproduces how often: (What percentage of the time does it reproduce?)

Every time I execute it.

• System information: (post the output of import pennylane as qml; qml.about())

master pennylane, recent catalyst

Source code and tracebacks

dev = qml.device('lightning.qubit', wires=2)

@qml.qjit
@qml.qnode(dev)
def circuit(phi):
    qml.Hadamard(0)
    qml.IsingXX(phi, wires=(0,1))
    return qml.density_matrix(wires=(0,1))

circuit(0.5)

array([0.68512454+0.j        , 0.        -0.17494102j,
       0.68512454+0.j        , 0.        -0.17494102j])

Please include any additional code snippets and error tracebacks related to the issue here.

Additional information

Any additional information, configuration or data that might be necessary to reproduce the issue.

[dime10]

Thanks, I don't expect this to work (at least it's not listed as a supported measurement process in the docs), but wouldn't that require a DM simulator?

[albi3ro]

We can calculate the density matrix from a state vector. We just can't do the reverse.

This may be an issue because DensityMatrixMP inherits from StateMP due to the fact that historically we treated a "density matrix" as just a StateMP with wires.

[dime10]

We can calculate the density matrix from a state vector. We just can't do the reverse.

This may be an issue because DensityMatrixMP inherits from StateMP due to the fact that historically we treated a "density matrix" as just a StateMP with wires.

Ohh I see, so this is a qml.state on a subset of wires 🤔 Yeah we probably convert StateMP to jaxpr with an isinstance check. We can probably ensure that this doesn't pass the verifier, or add support for this MP but that depends on what the device supports ...

[albi3ro]

Historically, density matrix was just state on a subset of wires. But we discovered that caused a lot of confusion and problems in pennylane, so we decided to separate them to the best of our ability.

This would be particularily relevant for catalyst, as a density matrix has a completely different shape and size than the state.

[josh146]

@lillian542 quick question: would your work on MP validation resolve this issue (since we would simply raise a better error message)?

[lillian542]

@josh146 this should raise an error now. My computer is currently rebuilding catalyst from the bottom up, so I can't check for sure, but I'll confirm that the example code raises a clear error on the latest catalyst branch as soon as my install is working again.

Assuming it does, can I close this issue?

[josh146]

Assuming it does, can I close this issue?

Yep!

[lillian542]

Unfortunately not quite solved yet, the current verification registers DensityMatrixMP as identical to StateMP, and allows it when it shouldn't. Basically what Christina said above:

This may be an issue because DensityMatrixMP inherits from StateMP due to the fact that historically we treated a "density matrix" as just a StateMP with wires.

is still an issue for DensityMatrixMP specifically.

[lillian542]

This now raises an NotImplementedError as discussed above, following https://github.com/PennyLaneAI/catalyst/pull/1118