Proof of concept constants are not hoisted.

[erick-xanadu]

Context: This is just a proof of concept, additional work is needed. What is this PR addressing? For this simple piece of code:

import numpy as np
import pennylane as qml
import jax

U = np.array([[-0.17111489, -0.69352236],
        [0.25053735, 0.60700543]])

@qml.qjit(target="mlir")
@qml.qnode(qml.device("lightning.qubit", wires=1))
def circuit():
    qml.QubitUnitary(U, wires=0)
    return qml.state()

print(circuit.mlir)

The MLIR produced always has the constant in the host program and not on the device program. I.e.,

builtin.module @circuit {
  func.func public @jit_circuit() -> tensor<2xcomplex<f64>> attributes {llvm.emit_c_interface} {
    // Constant is here
    %0 = "stablehlo.constant"() <{value = dense<[[-0.17111488999999999, -0.69352236], [0.25053734999999999, 0.60700542999999996]]> : tensor<2x2xf64>}> : () -> tensor<2x2xf64>

   // Passed as a parameter to the actual qnode
    %1 = "catalyst.launch_kernel"(%0) <{callee = @module_circuit::@circuit}> : (tensor<2x2xf64>) -> tensor<2xcomplex<f64>>
    func.return %1 : tensor<2xcomplex<f64>>
  }

I wondered how to get the constant inside the qnode. If we use jax.jit it is possible. However, there are other errors that stem from this decision, in particular dynamic shapes.

Other kind of errors also emerged, like:

• Can only differentiate quantum functions which do not call other quantum functions
• And that we have the invariant that quantum funcations cannot call other quantum functions.

As a way around this, I did outlining followed by inlining followed by outlining.

This is just a proof of concept intended more for discussion than for merging. It would be nice to have these constants inside the quantum function in MLIR to allow transformations to be easier.

With this proposal:

import numpy as np
import pennylane as qml
import jax


qml.capture.enable()

U = np.array([[-0.17111489, -0.69352236],
        [0.25053735, 0.60700543]])

@qml.qjit(target="mlir")
@qml.qnode(qml.device("lightning.qubit", wires=1))
def circuit():
    qml.QubitUnitary(U, wires=0)
    return qml.state()

qml.capture.disable()
print(circuit.mlir)

module @circuit {
  func.func public @jit_circuit() -> tensor<2xcomplex<f64>> attributes {llvm.emit_c_interface} {
    %0 = catalyst.launch_kernel @module_circuit::@circuit() : () -> tensor<2xcomplex<f64>>
    return %0 : tensor<2xcomplex<f64>>
  }
  module @module_circuit {
    func.func public @circuit() -> tensor<2xcomplex<f64>> attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
      %0 = call @"<unnamed wrapped function>"() : () -> tensor<2xcomplex<f64>>
      return %0 : tensor<2xcomplex<f64>>
    }

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[dime10]

Note that next quarter we are planning to switch over to the plxpr pipeline by default, so most frontend work should probably be targeting that pipeline from now on.

[erick-xanadu]

@dime10 , thanks! updated the important part.