Make Projector jit compatible (#5595)

**Description of the Change:**

Add some jit-friendly logic to `BasisStateProjector` and
`StateVectorProjector`.

**Benefits:**

Projector is now jit compatible.

We might finally be able to move forward with #4969 if we ever actually
had time to do that.

**Possible Drawbacks:**

Not testing trainability, as it doesn't really make sense.

Could potentially be made more efficient.

**Related GitHub Issues:**

Fixes #4977 [sc-52777]

---------

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Astral Cai <astral.cai@xanadu.ai>
Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>

doc/releases/changelog-dev.md

@@ -95,6 +95,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.Projector` is now compatible with jax-jit.
+ [(#5595)](https://github.com/PennyLaneAI/pennylane/pull/5595)
+
 * Finite shot circuits with a `qml.probs` measurement, both with a `wires` or `op` argument, can now be compiled with `jax.jit`.
  [(#5619)](https://github.com/PennyLaneAI/pennylane/pull/5619)
 

pennylane/math/single_dispatch.py

@@ -243,6 +243,7 @@ def _take_autograd(tensor, indices, axis=None):
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "isclose"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "atleast_1d"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "all"] = "tensorflow.experimental.numpy"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "vstack"] = "tensorflow.experimental.numpy"
 
 tf_fft_functions = [
  "fft",

pennylane/ops/op_math/composite.py

@@ -175,8 +175,9 @@ def eigvals(self):
  list(self.wires),
  )
  )
-
- return self._math_op(math.asarray(eigvals, like=math.get_deep_interface(eigvals)), axis=0)
+ framework = math.get_deep_interface(eigvals)
+ eigvals = [math.asarray(ei, like=framework) for ei in eigvals]
+ return self._math_op(math.vstack(eigvals), axis=0)
 
  @abc.abstractmethod
  def matrix(self, wire_order=None):

pennylane/ops/qubit/observables.py

@@ -435,14 +435,23 @@ class BasisStateProjector(Projector, Operation):
  r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where
  :math:`\phi` denotes a basis state."""
 
+ grad_method = None
+
  # The call signature should be the same as Projector.__new__ for the positional
  # arguments, but with free key word arguments.
  def __init__(self, state, wires, id=None):
  wires = Wires(wires)
- state = list(qml.math.toarray(state).astype(int))
 
- if not set(state).issubset({0, 1}):
- raise ValueError(f"Basis state must only consist of 0s and 1s; got {state}")
+ if qml.math.get_interface(state) == "jax":
+ dtype = qml.math.dtype(state)
+ if not (np.issubdtype(dtype, np.integer) or np.issubdtype(dtype, bool)):
+ raise ValueError("Basis state must consist of integers or booleans.")
+ else:
+ # state is index into data, rather than data, so cast it to built-ins when
+ # no need for tracing
+ state = tuple(qml.math.toarray(state).astype(int))
+ if not set(state).issubset({0, 1}):
+ raise ValueError(f"Basis state must only consist of 0s and 1s; got {state}")
 
  super().__init__(state, wires=wires, id=id)
 
@@ -471,7 +480,7 @@ def label(self, decimals=None, base_label=None, cache=None):
 
  if base_label is not None:
  return base_label
- basis_string = "".join(str(int(i)) for i in self.parameters[0])
+ basis_string = "".join(str(int(i)) for i in self.data[0])
  return f"|{basis_string}⟩⟨{basis_string}|"
 
  @staticmethod
@@ -497,7 +506,15 @@ def compute_matrix(basis_state): # pylint: disable=arguments-differ
  [0. 0. 0. 0.]
  [0. 0. 0. 0.]]
  """
- m = np.zeros((2 ** len(basis_state), 2 ** len(basis_state)))
+ shape = (2 ** len(basis_state), 2 ** len(basis_state))
+ if qml.math.get_interface(basis_state) == "jax":
+ idx = 0
+ for i, m in enumerate(basis_state):
+ idx = idx + (m << (len(basis_state) - i - 1))
+ mat = qml.math.zeros(shape, like=basis_state)
+ return mat.at[idx, idx].set(1.0)
+
+ m = np.zeros(shape)
  idx = int("".join(str(i) for i in basis_state), 2)
  m[idx, idx] = 1
  return m
@@ -528,6 +545,12 @@ def compute_eigvals(basis_state): # pylint: disable=arguments-differ
  >>> BasisStateProjector.compute_eigvals([0, 1])
  [0. 1. 0. 0.]
  """
+ if qml.math.get_interface(basis_state) == "jax":
+ idx = 0
+ for i, m in enumerate(basis_state):
+ idx = idx + (m << (len(basis_state) - i - 1))
+ eigvals = qml.math.zeros(2 ** len(basis_state), like=basis_state)
+ return eigvals.at[idx].set(1.0)
  w = np.zeros(2 ** len(basis_state))
  idx = int("".join(str(i) for i in basis_state), 2)
  w[idx] = 1
@@ -566,12 +589,12 @@ class StateVectorProjector(Projector):
  r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where
  :math:`\phi` denotes a state."""
 
+ grad_method = None
+
  # The call signature should be the same as Projector.__new__ for the positional
  # arguments, but with free key word arguments.
  def __init__(self, state, wires, id=None):
  wires = Wires(wires)
- state = list(qml.math.toarray(state))
-
  super().__init__(state, wires=wires, id=id)
 
  def __new__(cls, *_, **__): # pylint: disable=arguments-differ
@@ -680,9 +703,11 @@ def compute_eigvals(state_vector): # pylint: disable=arguments-differ,arguments
  >>> StateVectorProjector.compute_eigvals([0, 0, 1, 0])
  array([1, 0, 0, 0])
  """
- w = qml.math.zeros_like(state_vector)
+ precision = qml.math.get_dtype_name(state_vector)[-2:]
+ dtype = f"float{precision}" if precision in {"32", "64"} else "float64"
+ w = np.zeros(qml.math.shape(state_vector), dtype=dtype)
  w[0] = 1
- return w
+ return qml.math.convert_like(w, state_vector)
 
  @staticmethod
  def compute_diagonalizing_gates(
@@ -715,10 +740,27 @@ def compute_diagonalizing_gates(
  # Adapting the approach discussed in the link below to work with arbitrary complex-valued state vectors.
  # Alternatively, we could take the adjoint of the Mottonen decomposition for the state vector.
  # https://quantumcomputing.stackexchange.com/questions/10239/how-can-i-fill-a-unitary-knowing-only-its-first-column
- phase = qml.math.exp(-1j * qml.math.angle(state_vector[0]))
+
+ if qml.math.get_interface(state_vector) == "tensorflow":
+ dtype_name = qml.math.get_dtype_name(state_vector)
+ if dtype_name == "int32":
+ state_vector = qml.math.cast(state_vector, np.complex64)
+ elif dtype_name == "int64":
+ state_vector = qml.math.cast(state_vector, np.complex128)
+
+ angle = qml.math.angle(state_vector[0])
+ if qml.math.get_interface(angle) == "tensorflow":
+ if qml.math.get_dtype_name(angle) == "float32":
+ angle = qml.math.cast(angle, np.complex64)
+ else:
+ angle = qml.math.cast(angle, np.complex128)
+
+ phase = qml.math.exp(-1.0j * angle)
  psi = phase * state_vector
  denominator = qml.math.sqrt(2 + 2 * psi[0])
- psi = qml.math.set_index(psi, 0, psi[0] + 1) # psi[0] += 1, but JAX-JIT compatible
+ summed_array = np.zeros(qml.math.shape(psi), dtype=qml.math.get_dtype_name(psi))
+ summed_array[0] = 1.0
+ psi = psi + summed_array
  psi /= denominator
  u = 2 * qml.math.outer(psi, qml.math.conj(psi)) - qml.math.eye(len(psi))
  return [QubitUnitary(u, wires=wires)]

pennylane/utils.py

@@ -205,4 +205,4 @@ def expand_vector(vector, original_wires, expanded_wires):
  expanded_tensor, tuple(original_indices), tuple(wire_indices)
  )
 
- return qml.math.reshape(expanded_tensor, qudit_order**M)
+ return qml.math.reshape(expanded_tensor, (qudit_order**M,))

tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py

@@ -14,7 +14,7 @@
 """Integration tests for using the TensorFlow interface with a QNode"""
 import numpy as np
 
-# pylint: disable=too-many-arguments,too-few-public-methods,comparison-with-callable
+# pylint: disable=too-many-arguments,too-few-public-methods,comparison-with-callable, use-implicit-booleaness-not-comparison
 import pytest
 
 import pennylane as qml
@@ -955,8 +955,9 @@ def cost_fn(x, y):
  assert np.allclose(grad, expected, atol=tol, rtol=0)
 
  @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector
+ @pytest.mark.parametrize("dtype", ("int32", "int64"))
  def test_projector(
- self, state, dev, diff_method, grad_on_execution, device_vjp, tol, interface
+ self, state, dev, diff_method, grad_on_execution, device_vjp, tol, interface, dtype
  ):
  """Test that the variance of a projector is correctly returned"""
  kwargs = {
@@ -974,7 +975,7 @@ def test_projector(
  kwargs["num_directions"] = 20
  tol = TOL_FOR_SPSA
 
- P = tf.constant(state)
+ P = tf.constant(state, dtype=dtype)
 
  x, y = 0.765, -0.654
  weights = tf.Variable([x, y], dtype=tf.float64)

tests/ops/qubit/test_observables.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit tests for qubit observables."""
-# pylint: disable=protected-access
+# pylint: disable=protected-access, use-implicit-booleaness-not-comparison
 import functools
 import pickle
 
@@ -540,6 +540,44 @@ def test_serialization(self):
  assert qml.equal(new_proj, proj)
  assert new_proj.id == proj.id # Ensure they are identical
 
+ @pytest.mark.jax
+ def test_jit_measurement(self):
+ """Test that the measurement of a projector can be jitted."""
+
+ import jax
+
+ @jax.jit
+ @qml.qnode(qml.device("default.qubit"))
+ def circuit(state):
+ qml.X(1)
+ return qml.expval(qml.Projector(state, wires=(0, 1)))
+
+ state00 = jax.numpy.array([0, 0])
+ out00 = circuit(state00)
+ assert qml.math.allclose(out00, 0)
+ state01 = jax.numpy.array([0, 1])
+ out01 = circuit(state01)
+ assert qml.math.allclose(out01, 1)
+ state10 = jax.numpy.array([True, False])
+ out10 = circuit(state10)
+ assert qml.math.allclose(out10, 0)
+
+ with pytest.raises(ValueError, match=r"Basis state must consist of integers or booleans."):
+ circuit(jax.numpy.array([0.5, 0.6]))
+
+ @pytest.mark.jax
+ def test_jit_matrix(self):
+ """Test that computing the matrix of a projector is jittable."""
+
+ import jax
+
+ basis_state = jax.numpy.array([0, 1])
+ f = jax.jit(BasisStateProjector.compute_matrix)
+ out = f(basis_state)
+
+ expected = np.array([[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
+ assert qml.math.allclose(out, expected)
+
 
 class TestBasisStateProjector:
  """Tests for the basis state projector observable."""
@@ -716,6 +754,24 @@ def test_label_matrices_list_in_cache(self, projector):
  assert len(cache["matrices"]) == 2
  assert np.allclose(cache["matrices"][1], projector.parameters[0])
 
+ @pytest.mark.jax
+ def test_jit_execution(self):
+ """Test that executing a StateVectorProjector can be jitted."""
+
+ import jax
+
+ @jax.jit
+ @qml.qnode(qml.device("default.qubit"))
+ def circuit(state):
+ return qml.expval(qml.Projector(state, wires=(0, 1)))
+
+ basis_state = jax.numpy.array([1, 1, 1, 1.0]) / 2
+ out = circuit(basis_state)
+ assert qml.math.allclose(out, 0.25)
+
+ basis_state2 = jax.numpy.array([0, 0, 0, 0])
+ assert qml.math.allclose(circuit(basis_state2), 0)
+
 
 label_data = [
  (qml.Hermitian(np.eye(2), wires=1), "𝓗"),