paul0403
commented
10 months ago I have updated the four files, and I can build the updated catalyst with make runtime ENABLE_OPENQASM=ON and make frontend successfully. I am fairly certain what I did was correct: I will make a branch for the issue later.
The problem I have is with testing. I am trying to test with
import pennylane as qml
from catalyst import *
dev = qml.device("braket.local.qubit", wires=2, shots=100)
@qjit
@qml.qnode(dev)
def f(x: float):
qml.Hadamard(0)
qml.PSWAP(x, wires=[0, 1])
qml.ISWAP(wires=[1, 0])
return qml.probs()and see whether the f.jaxpr still does the decomposition, or does it just call the PSWAP and ISWAP instructions directly.
However, I'm running into an issue with the imports. I cannot seem to import jax on my machine:
>>> import pennylane
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/__init__.py", line 26, in <module>
import pennylane.numpy
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/numpy/__init__.py", line 87, in <module>
from .wrapper import wrap_arrays, extract_tensors, tensor_wrapper
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/numpy/wrapper.py", line 23, in <module>
from .tensor import tensor
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/numpy/tensor.py", line 25, in <module>
from pennylane.operation import Operator
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/operation.py", line 256, in <module>
from pennylane.math import expand_matrix
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/math/__init__.py", line 37, in <module>
from .matrix_manipulation import expand_matrix, reduce_matrices, get_batch_size
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/math/matrix_manipulation.py", line 24, in <module>
from pennylane.wires import Wires
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/wires.py", line 21, in <module>
from pennylane.pytrees import register_pytree
File "/home/paul/.local/lib/python3.10/site-packages/pennylane/pytrees.py", line 22, in <module>
import jax.tree_util as jax_tree_util
File "/home/paul/.local/lib/python3.10/site-packages/jax/__init__.py", line 35, in <module>
from jax import config as _config_module
File "/home/paul/.local/lib/python3.10/site-packages/jax/config.py", line 17, in <module>
from jax._src.config import config # noqa: F401
File "/home/paul/.local/lib/python3.10/site-packages/jax/_src/config.py", line 27, in <module>
from jax._src import lib
File "/home/paul/.local/lib/python3.10/site-packages/jax/_src/lib/__init__.py", line 84, in <module>
cpu_feature_guard.check_cpu_features()
RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.Moreover, importing catalyst is even worse:
>>> import catalyst
Illegal instruction (core dumped)I am working with Linux paul-Standard-PC-Q35-ICH9-2009 6.2.0-39-generic #40~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Thu Nov 16 10:53:04 UTC 2 x86_64 x86_64 x86_64 GNU/Linux (a UTM VM on an Intel i5 MAC). Since the environment is linux x86_64, jax should have no issues. I am a bit confused as to what happened.
The jax and jaxlib version are both 0.4.14, as confirmed by the make frontend script:
install Catalyst Frontend
/usr/bin/python3 -m pip install -e .
Defaulting to user installation because normal site-packages is not writeable
Obtaining file:///home/paul/catalyst
Installing build dependencies ... done
Checking if build backend supports build_editable ... done
Getting requirements to build editable ... done
Preparing editable metadata (pyproject.toml) ... done
Requirement already satisfied: pennylane>=0.32 in /home/paul/.local/lib/python3.10/site-packages (from PennyLane-Catalyst==0.3.3.dev0) (0.34.0)
Requirement already satisfied: jax==0.4.14 in /home/paul/.local/lib/python3.10/site-packages (from PennyLane-Catalyst==0.3.3.dev0) (0.4.14)
Requirement already satisfied: jaxlib==0.4.14 in /home/paul/.local/lib/python3.10/site-packages (from PennyLane-Catalyst==0.3.3.dev0) (0.4.14)
Requirement already satisfied: scipy in /home/paul/.local/lib/python3.10/site-packages (from PennyLane-Catalyst==0.3.3.dev0) (1.11.4)
Requirement already satisfied: tomlkit in /home/paul/.local/lib/python3.10/site-packages (from PennyLane-Catalyst==0.3.3.dev0) (0.12.3)
Requirement already satisfied: ml-dtypes>=0.2.0 in /home/paul/.local/lib/python3.10/site-packages (from jax==0.4.14->PennyLane-Catalyst==0.3.3.dev0) (0.3.2)
Requirement already satisfied: numpy>=1.22 in /home/paul/.local/lib/python3.10/site-packages (from jax==0.4.14->PennyLane-Catalyst==0.3.3.dev0) (1.26.3)
Requirement already satisfied: opt-einsum in /home/paul/.local/lib/python3.10/site-packages (from jax==0.4.14->PennyLane-Catalyst==0.3.3.dev0) (3.3.0)
Requirement already satisfied: networkx in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (3.2.1)
Requirement already satisfied: rustworkx in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (0.13.2)
Requirement already satisfied: autograd in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (1.6.2)
Requirement already satisfied: toml in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (0.10.2)
Requirement already satisfied: appdirs in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (1.4.4)
Requirement already satisfied: semantic-version>=2.7 in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (2.10.0)
Requirement already satisfied: autoray>=0.6.1 in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (0.6.7)
Requirement already satisfied: cachetools in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (5.3.2)
Requirement already satisfied: pennylane-lightning>=0.34 in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (0.34.0)
Requirement already satisfied: requests in /usr/lib/python3/dist-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (2.25.1)
Requirement already satisfied: typing-extensions in /home/paul/.local/lib/python3.10/site-packages (from pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (4.9.0)
Requirement already satisfied: future>=0.15.2 in /usr/lib/python3/dist-packages (from autograd->pennylane>=0.32->PennyLane-Catalyst==0.3.3.dev0) (0.18.2)
Building wheels for collected packages: PennyLane-Catalyst
Building editable for PennyLane-Catalyst (pyproject.toml) ... done
Created wheel for PennyLane-Catalyst: filename=PennyLane_Catalyst-0.3.3.dev0-0.editable-cp310-cp310-linux_x86_64.whl size=9512 sha256=edb2f5433622a14020cc6d0df4b1c90d38fbf21427c292b379e10a870f60faf3
Stored in directory: /tmp/pip-ephem-wheel-cache-mpd9e2ae/wheels/0f/1e/3c/b35aee31d4d6ae2980f5166ef6643177f0887a9c4baef6e6a0
Successfully built PennyLane-Catalyst
Installing collected packages: PennyLane-Catalyst
Attempting uninstall: PennyLane-Catalyst
Found existing installation: PennyLane-Catalyst 0.3.3.dev0
Uninstalling PennyLane-Catalyst-0.3.3.dev0:
Successfully uninstalled PennyLane-Catalyst-0.3.3.dev0
Successfully installed PennyLane-Catalyst-0.3.3.dev0
rm -r frontend/PennyLane_Catalyst.egg-infoAs a workaround I'm trying to build jax from source. Any help would be appreciated. Thanks.
josh146
maliasadi
Context
We would like to add support for the PennyLane PSWAP and ISWAP gates to our compilation architecture. Originally, no Catalyst device supported them natively, so they weren't considered, but the Braket plugin which was supported at a later point in time does come with support for them.
Goal
We would like to support the following example without having to rely on decomposition:
We can check that right decomposition occurs by inspecting the captured program:
PSWAPandISWAPare not present, instead we seeCNOT,SWAP, and single qubit gates.Technical details
Generally, Catalyst handles unsupported gates in the Python layer. For this, it relies on the tape expansion mechanism of the device chosen by the user, decomposing any gates not considered "supported". Additionally, Catalyst will also decompose any gates that are unsupported by the compilation pipeline, specifically the QIR interface between compiled program and runtime.
In order to support the new gates, they need to be added in several places:
the list of supported gates in the
QJITDeviceclass: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/frontend/catalyst/pennylane_extensions.py#L201a QIR-style function declaration for each gate needs to be added to the C-API of the runtime: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/runtime/include/RuntimeCAPI.h#L51-L52
an implementation of the C-API function, calling the
QuantumDevice'sNamedOperationfunction, similar to: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/runtime/lib/capi/RuntimeCAPI.cpp#L350-L355an implementation in the OpenQASM-based Braket backend device, by updating the operation map: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/runtime/lib/backend/openqasm/OpenQasmBuilder.hpp#L71-L74
Note that when installing Catalyst from source, the OpenQASM device is not build by default. Refer to the linked documentation for how to build and test the device.