diff --git a/.github/stable/all_interfaces.txt b/.github/stable/all_interfaces.txt index 2faae5935c5..4f9d3bb9ae4 100644 --- a/.github/stable/all_interfaces.txt +++ b/.github/stable/all_interfaces.txt @@ -5,48 +5,48 @@ astunparse==1.6.3 autograd==1.6.2 autoray==0.6.12 black==24.4.2 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cfgv==3.4.0 charset-normalizer==3.3.2 clarabel==0.9.0 click==8.1.7 contourpy==1.2.1 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 filelock==3.15.4 flaky==3.8.1 flatbuffers==24.3.25 -fonttools==4.53.0 -fsspec==2024.6.0 +fonttools==4.53.1 +fsspec==2024.6.1 future==1.0.0 -gast==0.5.4 +gast==0.6.0 google-pasta==0.2.0 -grpcio==1.64.1 +grpcio==1.65.1 h5py==3.11.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 -importlib_metadata==7.2.1 +importlib_metadata==8.1.0 importlib_resources==6.4.0 iniconfig==2.0.0 isort==5.13.2 jax==0.4.23 jaxlib==0.4.23 Jinja2==3.1.4 -keras==3.3.3 +keras==3.4.1 kiwisolver==1.4.5 lazy-object-proxy==1.10.0 libclang==18.1.1 Markdown==3.6 markdown-it-py==3.0.0 MarkupSafe==2.1.5 -matplotlib==3.9.0 +matplotlib==3.9.1 mccabe==0.6.1 mdurl==0.1.2 ml-dtypes==0.3.2 @@ -57,12 +57,12 @@ networkx==3.2.1 nodeenv==1.9.1 numpy==1.26.4 opt-einsum==3.3.0 -optree==0.11.0 -osqp==0.6.7 +optree==0.12.1 +osqp==0.6.7.post0 packaging==24.1 pathspec==0.12.1 -PennyLane_Lightning==0.37.0 -pillow==10.3.0 +PennyLane_Lightning==0.38.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -72,7 +72,7 @@ py-cpuinfo==9.0.0 Pygments==2.18.0 pylint==2.7.4 pyparsing==3.1.2 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -84,16 +84,16 @@ PyYAML==6.0.1 qdldl==0.1.7.post4 requests==2.32.3 rich==13.7.1 -rustworkx==0.14.2 +rustworkx==0.15.1 scipy==1.12.0 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 six==1.16.0 -sympy==1.12.1 +sympy==1.13.1 tensorboard==2.16.2 tensorboard-data-server==0.7.2 -tensorflow==2.16.1 -tensorflow-io-gcs-filesystem==0.37.0 +tensorflow==2.16.2 +tensorflow-io-gcs-filesystem==0.37.1 termcolor==2.4.0 tf_keras==2.16.0 toml==0.10.2 diff --git a/.github/stable/core.txt b/.github/stable/core.txt index f98bb878bae..80f61ac518c 100644 --- a/.github/stable/core.txt +++ b/.github/stable/core.txt @@ -3,43 +3,43 @@ astroid==2.6.6 autograd==1.6.2 autoray==0.6.12 black==24.4.2 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cfgv==3.4.0 charset-normalizer==3.3.2 clarabel==0.9.0 click==8.1.7 contourpy==1.2.1 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 filelock==3.15.4 flaky==3.8.1 -fonttools==4.53.0 +fonttools==4.53.1 future==1.0.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 importlib_resources==6.4.0 iniconfig==2.0.0 isort==5.13.2 kiwisolver==1.4.5 lazy-object-proxy==1.10.0 -matplotlib==3.9.0 +matplotlib==3.9.1 mccabe==0.6.1 mypy-extensions==1.0.0 networkx==3.2.1 nodeenv==1.9.1 numpy==1.26.4 -osqp==0.6.7 +osqp==0.6.7.post0 packaging==24.1 pathspec==0.12.1 -PennyLane_Lightning==0.37.0 -pillow==10.3.0 +PennyLane_Lightning==0.38.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -47,7 +47,7 @@ py==1.11.0 py-cpuinfo==9.0.0 pylint==2.7.4 pyparsing==3.1.2 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -58,9 +58,9 @@ python-dateutil==2.9.0.post0 PyYAML==6.0.1 qdldl==0.1.7.post4 requests==2.32.3 -rustworkx==0.14.2 +rustworkx==0.15.1 scipy==1.11.4 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 six==1.16.0 toml==0.10.2 diff --git a/.github/stable/doc.txt b/.github/stable/doc.txt index baf67d1e0a6..d9f67c83c97 100644 --- a/.github/stable/doc.txt +++ b/.github/stable/doc.txt @@ -9,8 +9,8 @@ attrs==23.2.0 autograd==1.6.2 autoray==0.6.12 Babel==2.15.0 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 charset-normalizer==3.3.2 cirq-core==1.3.0 contourpy==1.2.1 @@ -18,22 +18,22 @@ cycler==0.12.1 deprecation==2.1.0 docutils==0.16 duet==0.2.9 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 flatbuffers==24.3.25 -fonttools==4.53.0 +fonttools==4.53.1 frozenlist==1.4.1 -fsspec==2024.6.0 +fsspec==2024.6.1 future==1.0.0 gast==0.4.0 -google-auth==2.30.0 +google-auth==2.32.0 google-auth-oauthlib==0.4.6 google-pasta==0.2.0 graphviz==0.20.3 -grpcio==1.64.1 +grpcio==1.65.1 h5py==3.11.0 idna==3.7 imagesize==1.4.1 -importlib_metadata==7.2.1 +importlib_metadata==8.1.0 importlib_resources==6.4.0 iniconfig==2.0.0 jax==0.4.16 @@ -60,8 +60,8 @@ opt-einsum==3.3.0 packaging==24.1 pandas==2.2.2 pennylane-sphinx-theme==0.5.7 -PennyLane_Lightning==0.36.0 -pillow==10.3.0 +PennyLane_Lightning==0.37.0 +pillow==10.4.0 pluggy==1.5.0 protobuf==3.19.6 PubChemPy==1.0.4 @@ -73,7 +73,7 @@ Pygments==2.18.0 pygments-github-lexers==0.0.5 pyparsing==3.1.2 pyscf==2.6.2 -pytest==8.2.2 +pytest==8.3.1 python-dateutil==2.9.0.post0 pytz==2024.1 PyYAML==6.0.1 @@ -98,13 +98,13 @@ sphinxcontrib-jsmath==1.0.1 sphinxcontrib-qthelp==1.0.3 sphinxcontrib-serializinghtml==1.1.5 sphinxext-opengraph==0.6.3 -sympy==1.12.1 +sympy==1.13.1 tensorboard==2.11.2 tensorboard-data-server==0.6.1 tensorboard-plugin-wit==1.8.1 tensorflow==2.11.1 tensorflow-estimator==2.11.0 -tensorflow-io-gcs-filesystem==0.37.0 +tensorflow-io-gcs-filesystem==0.37.1 tensornetwork==0.3.0 termcolor==2.4.0 toml==0.10.2 diff --git a/.github/stable/external.txt b/.github/stable/external.txt index 23201352049..507721a88ae 100644 --- a/.github/stable/external.txt +++ b/.github/stable/external.txt @@ -15,8 +15,8 @@ Babel==2.15.0 beautifulsoup4==4.12.3 black==24.4.2 bleach==6.1.0 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cffi==1.16.0 cfgv==3.4.0 charset-normalizer==3.3.2 @@ -25,43 +25,43 @@ click==8.1.7 comm==0.2.2 contourpy==1.2.1 cotengra==0.6.2 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 cytoolz==0.12.3 -debugpy==1.8.1 +debugpy==1.8.2 decorator==5.1.1 defusedxml==0.7.1 -diastatic-malt==2.15.1 +diastatic-malt==2.15.2 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 executing==2.0.1 fastjsonschema==2.20.0 filelock==3.15.4 flaky==3.8.1 flatbuffers==24.3.25 -fonttools==4.53.0 +fonttools==4.53.1 fqdn==1.5.1 future==1.0.0 -gast==0.5.4 +gast==0.6.0 google-pasta==0.2.0 -grpcio==1.64.1 +grpcio==1.65.1 h11==0.14.0 h5py==3.11.0 httpcore==1.0.5 httpx==0.27.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 -importlib_metadata==7.2.1 +importlib_metadata==8.1.0 importlib_resources==6.4.0 iniconfig==2.0.0 -ipykernel==6.29.4 +ipykernel==6.29.5 ipython==8.18.1 ipython-genutils==0.2.0 -ipywidgets==7.8.1 +ipywidgets==7.8.2 isoduration==20.11.0 isort==5.13.2 jax==0.4.23 @@ -70,19 +70,19 @@ jedi==0.19.1 Jinja2==3.1.4 json5==0.9.25 jsonpointer==3.0.0 -jsonschema==4.22.0 +jsonschema==4.23.0 jsonschema-specifications==2023.12.1 jupyter-events==0.10.0 jupyter-lsp==2.2.5 jupyter_client==8.6.2 jupyter_core==5.7.2 -jupyter_server==2.14.1 +jupyter_server==2.14.2 jupyter_server_terminals==0.5.3 -jupyterlab==4.2.2 -jupyterlab-widgets==1.1.7 +jupyterlab==4.2.4 jupyterlab_pygments==0.3.0 -jupyterlab_server==2.27.2 -keras==3.3.3 +jupyterlab_server==2.27.3 +jupyterlab_widgets==1.1.8 +keras==3.4.1 kiwisolver==1.4.5 lark==1.1.9 lazy-object-proxy==1.10.0 @@ -91,7 +91,7 @@ llvmlite==0.43.0 Markdown==3.6 markdown-it-py==3.0.0 MarkupSafe==2.1.5 -matplotlib==3.9.0 +matplotlib==3.9.1 matplotlib-inline==0.1.7 mccabe==0.6.1 mdurl==0.1.2 @@ -110,8 +110,8 @@ notebook_shim==0.2.4 numba==0.60.0 numpy==1.26.4 opt-einsum==3.3.0 -optree==0.11.0 -osqp==0.6.7 +optree==0.12.1 +osqp==0.6.7.post0 overrides==7.7.0 packaging==24.1 pandocfilters==1.5.1 @@ -120,7 +120,7 @@ pathspec==0.12.1 PennyLane-Catalyst==0.7.0 PennyLane_Lightning==0.37.0 pexpect==4.9.0 -pillow==10.3.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -129,7 +129,7 @@ prompt_toolkit==3.0.47 protobuf==4.25.3 psutil==6.0.0 ptyprocess==0.7.0 -pure-eval==0.2.2 +pure_eval==0.2.3 py==1.11.0 py-cpuinfo==9.0.0 pycparser==2.22 @@ -137,7 +137,7 @@ Pygments==2.18.0 pylint==2.7.4 pyparsing==3.1.2 pyperclip==1.9.0 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -149,16 +149,16 @@ PyYAML==6.0.1 pyzmq==26.0.3 pyzx==0.8.0 qdldl==0.1.7.post4 -quimb==1.8.2 +quimb==1.8.4 referencing==0.35.1 requests==2.32.3 rfc3339-validator==0.1.4 rfc3986-validator==0.1.1 rich==13.7.1 -rpds-py==0.18.1 -rustworkx==0.14.2 +rpds-py==0.19.0 +rustworkx==0.15.1 scipy==1.12.0 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 Send2Trash==1.8.3 six==1.16.0 @@ -168,15 +168,15 @@ stack-data==0.6.3 stim==1.13.0 tensorboard==2.16.2 tensorboard-data-server==0.7.2 -tensorflow==2.16.1 -tensorflow-io-gcs-filesystem==0.37.0 +tensorflow==2.16.2 +tensorflow-io-gcs-filesystem==0.37.1 termcolor==2.4.0 terminado==0.18.1 tf_keras==2.16.0 tinycss2==1.3.0 toml==0.10.2 tomli==2.0.1 -tomlkit==0.12.5 +tomlkit==0.13.0 toolz==0.12.1 tornado==6.4.1 tqdm==4.66.4 @@ -191,6 +191,6 @@ webcolors==24.6.0 webencodings==0.5.1 websocket-client==1.8.0 Werkzeug==3.0.3 -widgetsnbextension==3.6.6 +widgetsnbextension==3.6.7 wrapt==1.12.1 zipp==3.19.2 diff --git a/.github/stable/jax.txt b/.github/stable/jax.txt index f0a062792e4..23134473c21 100644 --- a/.github/stable/jax.txt +++ b/.github/stable/jax.txt @@ -3,28 +3,28 @@ astroid==2.6.6 autograd==1.6.2 autoray==0.6.12 black==24.4.2 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cfgv==3.4.0 charset-normalizer==3.3.2 clarabel==0.9.0 click==8.1.7 contourpy==1.2.1 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 filelock==3.15.4 flaky==3.8.1 -fonttools==4.53.0 +fonttools==4.53.1 future==1.0.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 -importlib_metadata==7.2.1 +importlib_metadata==8.1.0 importlib_resources==6.4.0 iniconfig==2.0.0 isort==5.13.2 @@ -32,7 +32,7 @@ jax==0.4.23 jaxlib==0.4.23 kiwisolver==1.4.5 lazy-object-proxy==1.10.0 -matplotlib==3.9.0 +matplotlib==3.9.1 mccabe==0.6.1 ml-dtypes==0.4.0 mypy-extensions==1.0.0 @@ -40,11 +40,11 @@ networkx==3.2.1 nodeenv==1.9.1 numpy==1.26.4 opt-einsum==3.3.0 -osqp==0.6.7 +osqp==0.6.7.post0 packaging==24.1 pathspec==0.12.1 -PennyLane_Lightning==0.37.0 -pillow==10.3.0 +PennyLane_Lightning==0.38.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -52,7 +52,7 @@ py==1.11.0 py-cpuinfo==9.0.0 pylint==2.7.4 pyparsing==3.1.2 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -63,9 +63,9 @@ python-dateutil==2.9.0.post0 PyYAML==6.0.1 qdldl==0.1.7.post4 requests==2.32.3 -rustworkx==0.14.2 +rustworkx==0.15.1 scipy==1.12.0 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 six==1.16.0 toml==0.10.2 diff --git a/.github/stable/tf.txt b/.github/stable/tf.txt index a2df434aafc..fa9d357bbef 100644 --- a/.github/stable/tf.txt +++ b/.github/stable/tf.txt @@ -5,44 +5,44 @@ astunparse==1.6.3 autograd==1.6.2 autoray==0.6.12 black==24.4.2 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cfgv==3.4.0 charset-normalizer==3.3.2 clarabel==0.9.0 click==8.1.7 contourpy==1.2.1 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 filelock==3.15.4 flaky==3.8.1 flatbuffers==24.3.25 -fonttools==4.53.0 +fonttools==4.53.1 future==1.0.0 -gast==0.5.4 +gast==0.6.0 google-pasta==0.2.0 -grpcio==1.64.1 +grpcio==1.65.1 h5py==3.11.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 -importlib_metadata==7.2.1 +importlib_metadata==8.1.0 importlib_resources==6.4.0 iniconfig==2.0.0 isort==5.13.2 -keras==3.3.3 +keras==3.4.1 kiwisolver==1.4.5 lazy-object-proxy==1.10.0 libclang==18.1.1 Markdown==3.6 markdown-it-py==3.0.0 MarkupSafe==2.1.5 -matplotlib==3.9.0 +matplotlib==3.9.1 mccabe==0.6.1 mdurl==0.1.2 ml-dtypes==0.3.2 @@ -52,12 +52,12 @@ networkx==3.2.1 nodeenv==1.9.1 numpy==1.26.4 opt-einsum==3.3.0 -optree==0.11.0 -osqp==0.6.7 +optree==0.12.1 +osqp==0.6.7.post0 packaging==24.1 pathspec==0.12.1 -PennyLane_Lightning==0.37.0 -pillow==10.3.0 +PennyLane_Lightning==0.38.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -67,7 +67,7 @@ py-cpuinfo==9.0.0 Pygments==2.18.0 pylint==2.7.4 pyparsing==3.1.2 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -79,15 +79,15 @@ PyYAML==6.0.1 qdldl==0.1.7.post4 requests==2.32.3 rich==13.7.1 -rustworkx==0.14.2 +rustworkx==0.15.1 scipy==1.11.4 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 six==1.16.0 tensorboard==2.16.2 tensorboard-data-server==0.7.2 -tensorflow==2.16.1 -tensorflow-io-gcs-filesystem==0.37.0 +tensorflow==2.16.2 +tensorflow-io-gcs-filesystem==0.37.1 termcolor==2.4.0 tf_keras==2.16.0 toml==0.10.2 diff --git a/.github/stable/torch.txt b/.github/stable/torch.txt index b8afd4e773d..ea0d1dbb955 100644 --- a/.github/stable/torch.txt +++ b/.github/stable/torch.txt @@ -3,27 +3,27 @@ astroid==2.6.6 autograd==1.6.2 autoray==0.6.12 black==24.4.2 -cachetools==5.3.3 -certifi==2024.6.2 +cachetools==5.4.0 +certifi==2024.7.4 cfgv==3.4.0 charset-normalizer==3.3.2 clarabel==0.9.0 click==8.1.7 contourpy==1.2.1 -coverage==7.5.4 +coverage==7.6.0 cvxopt==1.3.2 cvxpy==1.5.2 cycler==0.12.1 distlib==0.3.8 ecos==2.0.14 -exceptiongroup==1.2.1 +exceptiongroup==1.2.2 execnet==2.1.1 filelock==3.15.4 flaky==3.8.1 -fonttools==4.53.0 -fsspec==2024.6.0 +fonttools==4.53.1 +fsspec==2024.6.1 future==1.0.0 -identify==2.5.36 +identify==2.6.0 idna==3.7 importlib_resources==6.4.0 iniconfig==2.0.0 @@ -32,18 +32,18 @@ Jinja2==3.1.4 kiwisolver==1.4.5 lazy-object-proxy==1.10.0 MarkupSafe==2.1.5 -matplotlib==3.9.0 +matplotlib==3.9.1 mccabe==0.6.1 mpmath==1.3.0 mypy-extensions==1.0.0 networkx==3.2.1 nodeenv==1.9.1 numpy==1.26.4 -osqp==0.6.7 +osqp==0.6.7.post0 packaging==24.1 pathspec==0.12.1 -PennyLane_Lightning==0.37.0 -pillow==10.3.0 +PennyLane_Lightning==0.38.0 +pillow==10.4.0 platformdirs==4.2.2 pluggy==1.5.0 pre-commit==3.7.1 @@ -51,7 +51,7 @@ py==1.11.0 py-cpuinfo==9.0.0 pylint==2.7.4 pyparsing==3.1.2 -pytest==8.2.2 +pytest==8.3.1 pytest-benchmark==4.0.0 pytest-cov==5.0.0 pytest-forked==1.6.0 @@ -62,12 +62,12 @@ pytorch-triton-rocm==2.3.0 PyYAML==6.0.1 qdldl==0.1.7.post4 requests==2.32.3 -rustworkx==0.14.2 +rustworkx==0.15.1 scipy==1.11.4 -scs==3.2.4.post3 +scs==3.2.6 semantic-version==2.10.0 six==1.16.0 -sympy==1.12.1 +sympy==1.13.1 toml==0.10.2 tomli==2.0.1 torch==2.3.0+rocm6.0 diff --git a/.github/workflows/core_tests_durations.json b/.github/workflows/core_tests_durations.json index bd00edf8f6f..7558d3087e6 100644 --- a/.github/workflows/core_tests_durations.json +++ b/.github/workflows/core_tests_durations.json @@ -13182,7 +13182,6 @@ "ops/op_math/test_composite.py::TestConstruction::test_map_wires[True-expected_overlapping_ops1]": 0.001044666045345366, "ops/op_math/test_composite.py::TestConstruction::test_ndim_params_raises_error": 0.0008577080443501472, "ops/op_math/test_composite.py::TestConstruction::test_parameters": 0.0009155830484814942, - "ops/op_math/test_composite.py::TestConstruction::test_queue_idx": 0.0008379160426557064, "ops/op_math/test_composite.py::TestConstruction::test_raise_error_fewer_than_2_operands": 0.0010355000267736614, "ops/op_math/test_composite.py::TestConstruction::test_tensor_and_hamiltonian_converted": 0.0014357499894686043, "ops/op_math/test_composite.py::TestMscMethods::test_copy[ops_lst0]": 0.0009615410235710442, diff --git a/doc/code/qml.rst b/doc/code/qml.rst index fa32159fefb..c568ea5ab26 100644 --- a/doc/code/qml.rst +++ b/doc/code/qml.rst @@ -6,4 +6,4 @@ qml .. automodapi:: pennylane :no-heading: :include-all-objects: - :skip: Version, SimpleSpec, plugin_devices, plugin_converters, default_config, reload, version_info, defaultdict + :skip: plugin_converters, default_config, version_info, defaultdict diff --git a/doc/code/qml_tape.rst b/doc/code/qml_tape.rst index 951a3a69c39..c83a418b5b9 100644 --- a/doc/code/qml_tape.rst +++ b/doc/code/qml_tape.rst @@ -67,4 +67,5 @@ and a reduction in unintended side effects, ``QuantumScript`` is strictly used i .. automodapi:: pennylane.tape :no-main-docstr: :include-all-objects: + :skip: QuantumTapeBatch :inheritance-diagram: \ No newline at end of file diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst index 768527a7079..7f469e13359 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -9,6 +9,64 @@ deprecations are listed below. Pending deprecations -------------------- +* All of the legacy devices (any with the name ``default.qubit.{autograd,torch,tf,jax,legacy}``) are deprecated. Use ``default.qubit`` instead, + as it supports backpropagation for the many backends the legacy devices support. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``decomp_depth`` argument in ``qml.device`` is deprecated. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``max_expansion`` argument in ``qml.QNode`` is deprecated. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The functions ``qml.transforms.sum_expand`` and ``qml.transforms.hamiltonian_expand`` are deprecated. + Instead, ``qml.transforms.split_non_commuting`` can be used for equivalent behaviour. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``expansion_strategy`` attribute of ``qml.QNode`` is deprecated. + Users should make use of ``qml.workflow.construct_batch``, should they require fine control over the output tape(s). + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``expansion_strategy`` argument in ``qml.specs``, ``qml.draw``, and ``qml.draw_mpl`` is deprecated. + Instead, use the ``level`` argument which provides a superset of options. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``expand_fn`` argument in ``qml.execute`` is deprecated. + Instead, please create a ``qml.transforms.core.TransformProgram`` with the desired preprocessing and pass it to the ``transform_program`` argument of ``qml.execute``. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``max_expansion`` argument in ``qml.execute`` is deprecated. + Instead, please use ``qml.devices.preprocess.decompose`` with the desired expansion level, add it to a ``TransformProgram``, and pass it to the ``transform_program`` argument of ``qml.execute``. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``override_shots`` argument in ``qml.execute`` is deprecated. + Instead, please add the shots to the ``QuantumTape``\ s to be executed. + + - Deprecated in v0.38 + - Will be removed in v0.39 + +* The ``device_batch_transform`` argument in ``qml.execute`` is deprecated. + Instead, please create a ``qml.transforms.core.TransformProgram`` with the desired preprocessing and pass it to the ``transform_program`` argument of ``qml.execute``. + + - Deprecated in v0.38 + - Will be removed in v0.39 + * The functions ``qml.qinfo.classical_fisher`` and ``qml.qinfo.quantum_fisher`` are deprecated since they are being migrated to the ``qml.gradients`` module. Therefore, ``qml.gradients.classical_fisher`` and ``qml.gradients.quantum_fisher`` should be used instead. @@ -65,6 +123,10 @@ Other deprecations Completed deprecation cycles ---------------------------- +* ``queue_idx`` attribute has been removed from the ``Operator``, ``CompositeOp``, and ``SymboliOp`` classes. Instead, the index is now stored as the label of the ``CircuitGraph.graph`` nodes. + + - Deprecated in v0.38 + - Removed in v0.38 * ``qml.from_qasm`` no longer removes measurements from the QASM code. Use ``measurements=[]`` to remove measurements from the original circuit. diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 1622d85ff57..863c0da6f48 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -13,6 +13,7 @@ * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear combination of unitaries. [(#5756)](https://github.com/PennyLaneAI/pennylane/pull/5756) + [(#5987)](https://github.com/PennyLaneAI/pennylane/pull/5987) * The `split_to_single_terms` transform is added. This transform splits expectation values of sums into multiple single-term measurements on a single tape, providing better support for simulators @@ -22,8 +23,19 @@ * `SProd.terms` now flattens out the terms if the base is a multi-term observable. [(#5885)](https://github.com/PennyLaneAI/pennylane/pull/5885) +* A new method `to_mat` has been added to the `FermiWord` and `FermiSentence` classes, which allows + computing the matrix representation of these Fermi operators. + [(#5920)](https://github.com/PennyLaneAI/pennylane/pull/5920) + +* New functionality has been added to natively support exponential extrapolation when using the `mitigate_with_zne`. This allows + users to have more control over the error mitigation protocol without needing to add further dependencies. + [(#5972)](https://github.com/PennyLaneAI/pennylane/pull/5972) +

Improvements πŸ› 

+* Added the `compute_sparse_matrix` method for `qml.ops.qubit.BasisStateProjector`. + [(#5790)](https://github.com/PennyLaneAI/pennylane/pull/5790) + * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations [(#5995)](https://github.com/PennyLaneAI/pennylane/pull/5995) @@ -50,6 +62,55 @@ * Molecules and Hamiltonians can now be constructed for all the elements present in the periodic table. [(#5821)](https://github.com/PennyLaneAI/pennylane/pull/5821) +* `qml.for_loop` and `qml.while_loop` now fallback to standard Python control + flow if `@qjit` is not present, allowing the same code to work with and without + `@qjit` without any rewrites. + [(#6014)](https://github.com/PennyLaneAI/pennylane/pull/6014) + + ```python + dev = qml.device("lightning.qubit", wires=3) + + @qml.qnode(dev) + def circuit(x, n): + + @qml.for_loop(0, n, 1) + def init_state(i): + qml.Hadamard(wires=i) + + init_state() + + @qml.for_loop(0, n, 1) + def apply_operations(i, x): + qml.RX(x, wires=i) + + @qml.for_loop(i + 1, n, 1) + def inner(j): + qml.CRY(x**2, [i, j]) + + inner() + return jnp.sin(x) + + apply_operations(x) + return qml.probs() + ``` + + ```pycon + >>> print(qml.draw(circuit)(0.5, 3)) + 0: ──H──RX(0.50)─╭●────────╭●─────────────────────────────────────── Probs + 1: ──H───────────╰RY(0.25)─│──────────RX(0.48)─╭●─────────────────── Probs + 2: ──H─────────────────────╰RY(0.25)───────────╰RY(0.23)──RX(0.46)── Probs + >>> circuit(0.5, 3) + array([0.125 , 0.125 , 0.09949758, 0.15050242, 0.07594666, + 0.11917543, 0.08942104, 0.21545687]) + >>> qml.qjit(circuit)(0.5, 3) + Array([0.125 , 0.125 , 0.09949758, 0.15050242, 0.07594666, + 0.11917543, 0.08942104, 0.21545687], dtype=float64) + ``` + +* The `qubit_observable` function is modified to return an ascending wire order for molecular + Hamiltonians. + [(#5950)](https://github.com/PennyLaneAI/pennylane/pull/5950) +

Community contributions πŸ₯³

* `DefaultQutritMixed` readout error has been added using parameters `readout_relaxation_probs` and @@ -60,11 +121,19 @@

Breaking changes πŸ’”

-* ``qml.from_qasm`` no longer removes measurements from the QASM code. Use - ``measurements=[]`` to remove measurements from the original circuit. +* The `CircuitGraph.graph` rustworkx graph now stores indices into the circuit as the node labels, + instead of the operator/ measurement itself. This allows the same operator to occur multiple times in + the circuit. + [(#5907)](https://github.com/PennyLaneAI/pennylane/pull/5907) + +* `queue_idx` attribute has been removed from the `Operator`, `CompositeOp`, and `SymboliOp` classes. + [(#6005)](https://github.com/PennyLaneAI/pennylane/pull/6005) + +* `qml.from_qasm` no longer removes measurements from the QASM code. Use + `measurements=[]` to remove measurements from the original circuit. [(#5982)](https://github.com/PennyLaneAI/pennylane/pull/5982) - -* ``qml.transforms.map_batch_transform`` has been removed, since transforms can be applied directly to a batch of tapes. + +* `qml.transforms.map_batch_transform` has been removed, since transforms can be applied directly to a batch of tapes. See :func:`~.pennylane.transform` for more information. [(#5981)](https://github.com/PennyLaneAI/pennylane/pull/5981) @@ -73,10 +142,51 @@

Deprecations πŸ‘‹

+* The `decomp_depth` argument in `qml.device` has been deprecated. + [(#6026)](https://github.com/PennyLaneAI/pennylane/pull/6026) + +* The `max_expansion` argument in `qml.QNode` has been deprecated. + [(#6026)](https://github.com/PennyLaneAI/pennylane/pull/6026) + +* The `expansion_strategy` attribute in the `QNode` class is deprecated. + [(#5989)](https://github.com/PennyLaneAI/pennylane/pull/5989) + +* The `expansion_strategy` argument has been deprecated in all of `qml.draw`, `qml.draw_mpl`, and `qml.specs`. + The `level` argument should be used instead. + [(#5989)](https://github.com/PennyLaneAI/pennylane/pull/5989) + +* `Operator.expand` has been deprecated. Users should simply use `qml.tape.QuantumScript(op.decomposition())` + for equivalent behaviour. + [(#5994)](https://github.com/PennyLaneAI/pennylane/pull/5994) + +* `pennylane.transforms.sum_expand` and `pennylane.transforms.hamiltonian_expand` have been deprecated. + Users should instead use `pennylane.transforms.split_non_commuting` for equivalent behaviour. + [(#6003)](https://github.com/PennyLaneAI/pennylane/pull/6003) + +* The `expand_fn` argument in `qml.execute` has been deprecated. + Instead, please create a `qml.transforms.core.TransformProgram` with the desired preprocessing and pass it to the `transform_program` argument of `qml.execute`. + [(#5984)](https://github.com/PennyLaneAI/pennylane/pull/5984) + +* The `max_expansion` argument in `qml.execute` has been deprecated. + Instead, please use `qml.devices.preprocess.decompose` with the desired expansion level, add it to a `TransformProgram` and pass it to the `transform_program` argument of `qml.execute`. + [(#5984)](https://github.com/PennyLaneAI/pennylane/pull/5984) + +* The `override_shots` argument in `qml.execute` is deprecated. + Instead, please add the shots to the `QuantumTape`'s to be executed. + [(#5984)](https://github.com/PennyLaneAI/pennylane/pull/5984) + +* The `device_batch_transform` argument in `qml.execute` is deprecated. + Instead, please create a `qml.transforms.core.TransformProgram` with the desired preprocessing and pass it to the `transform_program` argument of `qml.execute`. + [(#5984)](https://github.com/PennyLaneAI/pennylane/pull/5984) + * `pennylane.qinfo.classical_fisher` and `pennylane.qinfo.quantum_fisher` have been deprecated. Instead, use `pennylane.gradients.classical_fisher` and `pennylane.gradients.quantum_fisher`. [(#5985)](https://github.com/PennyLaneAI/pennylane/pull/5985) +* The legacy devices `default.qubit.{autograd,torch,tf,jax,legacy}` are deprecated. + Instead, use `default.qubit` as it now supports backpropagation through the several backends. + [(#5997)](https://github.com/PennyLaneAI/pennylane/pull/5997) +

Documentation πŸ“

* Improves the docstring for `QuantumScript.expand` and `qml.tape.tape.expand_tape`. @@ -84,6 +194,15 @@

Bug fixes πŸ›

+* Fixed a bug in `qml.SPSAOptimizer` that ignored keyword arguments in the objective function. + [(#6027)](https://github.com/PennyLaneAI/pennylane/pull/6027) + +* `dynamic_one_shot` was broken for old-API devices since `override_shots` was deprecated. + [(#6024)](https://github.com/PennyLaneAI/pennylane/pull/6024) + +* `CircuitGraph` can now handle circuits with the same operation instance occuring multiple times. + [(#5907)](https://github.com/PennyLaneAI/pennylane/pull/5907) + * `qml.QSVT` is updated to store wire order correctly. [(#5959)](https://github.com/PennyLaneAI/pennylane/pull/5959) @@ -94,6 +213,7 @@ * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types. [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969) +

Contributors ✍️

This release contains contributions from (in alphabetical order): @@ -106,9 +226,14 @@ Ahmed Darwish, Lillian M. A. Frederiksen, Pietropaolo Frisoni, Emiliano Godinez, +Renke Huang, +Josh Izaac, +Soran Jahangiri, Christina Lee, Austin Huang, William Maxwell, Vincent Michaud-Rioux, +Anurav Modak, Mudit Pandey, -Erik Schultheis. +Erik Schultheis, +nate stemen. diff --git a/pennylane/__init__.py b/pennylane/__init__.py index 04b7e4f24f9..8bc01c74fa1 100644 --- a/pennylane/__init__.py +++ b/pennylane/__init__.py @@ -15,12 +15,9 @@ This is the top level module from which all basic functions and classes of PennyLane can be directly imported. """ -from importlib import reload, metadata -from sys import version_info import numpy as _np -from semantic_version import SimpleSpec, Version from pennylane.boolean_fn import BooleanFn import pennylane.numpy @@ -154,6 +151,8 @@ import pennylane.noise from pennylane.noise import NoiseModel +from pennylane.devices.device_constructor import device, refresh_devices + # Look for an existing configuration file default_config = Configuration("config.toml") @@ -175,262 +174,10 @@ def __getattr__(name): return pennylane.ops.LinearCombination return pennylane.ops.Hamiltonian - raise AttributeError(f"module 'pennylane' has no attribute '{name}'") - - -def _get_device_entrypoints(): - """Returns a dictionary mapping the device short name to the - loadable entrypoint""" - entries = ( - metadata.entry_points()["pennylane.plugins"] - if version_info[:2] == (3, 9) - # pylint:disable=unexpected-keyword-arg - else metadata.entry_points(group="pennylane.plugins") - ) - return {entry.name: entry for entry in entries} - - -def refresh_devices(): - """Scan installed PennyLane plugins to refresh the device list.""" - - # This function does not return anything; instead, it has a side effect - # which is to update the global plugin_devices variable. - - # We wish to retain the behaviour of a global plugin_devices dictionary, - # as re-importing metadata can be a very slow operation on systems - # with a large number of installed packages. - global plugin_devices # pylint:disable=global-statement - - reload(metadata) - plugin_devices = _get_device_entrypoints() - - -# get list of installed devices -plugin_devices = _get_device_entrypoints() - - -# pylint: disable=protected-access -def device(name, *args, **kwargs): - r""" - Load a device and return the instance. - - This function is used to load a particular quantum device, - which can then be used to construct QNodes. - - PennyLane comes with support for the following devices: - - * :mod:`'default.qubit' `: a simple - state simulator of qubit-based quantum circuit architectures. - - * :mod:`'default.mixed' `: a mixed-state - simulator of qubit-based quantum circuit architectures. - - * ``'lightning.qubit'``: a more performant state simulator of qubit-based - quantum circuit architectures written in C++. - - * :mod:`'default.qutrit' `: a simple - state simulator of qutrit-based quantum circuit architectures. - - * :mod:`'default.qutrit.mixed' `: a - mixed-state simulator of qutrit-based quantum circuit architectures. - - * :mod:`'default.gaussian' `: a simple simulator - of Gaussian states and operations on continuous-variable circuit architectures. - - * :mod:`'default.clifford' `: an efficient - simulator of Clifford circuits. - - * :mod:`'default.tensor' `: a simulator - of quantum circuits based on tensor networks. - - Additional devices are supported through plugins β€” see - the `available plugins `_ for more - details. To list all currently installed devices, run - :func:`qml.about `. - - Args: - name (str): the name of the device to load - wires (int): the number of wires (subsystems) to initialise - the device with. Note that this is optional for certain - devices, such as ``default.qubit`` - - Keyword Args: - config (pennylane.Configuration): a PennyLane configuration object - that contains global and/or device specific configurations. - custom_decomps (Dict[Union(str, Operator), Callable]): Custom - decompositions to be applied by the device at runtime. - decomp_depth (int): For when custom decompositions are specified, - the maximum expansion depth used by the expansion function. - - All devices must be loaded by specifying their **short-name** as listed above, - followed by the **wires** (subsystems) you wish to initialize. The ``wires`` - argument can be an integer, in which case the wires of the device are addressed - by consecutive integers: - - .. code-block:: python - - dev = qml.device('default.qubit', wires=5) - - def circuit(): - qml.Hadamard(wires=1) - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[3, 4]) - ... - - The ``wires`` argument can also be a sequence of unique numbers or strings, specifying custom wire labels - that the user employs to address the wires: - - .. code-block:: python + if name == "plugin_devices": + return pennylane.devices.device_constructor.plugin_devices - dev = qml.device('default.qubit', wires=['ancilla', 'q11', 'q12', -1, 1]) - - def circuit(): - qml.Hadamard(wires='q11') - qml.Hadamard(wires=['ancilla']) - qml.CNOT(wires=['q12', -1]) - ... - - On some newer devices, such as ``default.qubit``, the ``wires`` argument can be omitted altogether, - and instead the wires will be computed when executing a circuit depending on its contents. - - >>> dev = qml.device("default.qubit") - - Most devices accept a ``shots`` argument which specifies how many circuit executions - are used to estimate stochastic return values. As an example, ``qml.sample()`` measurements - will return as many samples as specified in the shots argument. The shots argument can be - changed on a per-call basis using the built-in ``shots`` keyword argument. Note that the - ``shots`` argument can be a single integer or a list of shot values. - - .. code-block:: python - - dev = qml.device('default.qubit', wires=1, shots=10) - - @qml.qnode(dev) - def circuit(a): - qml.RX(a, wires=0) - return qml.sample(qml.Z(0)) - - >>> circuit(0.8) # 10 samples are returned - array([ 1, 1, 1, 1, -1, 1, 1, -1, 1, 1]) - >>> circuit(0.8, shots=[3, 4, 4]) # default is overwritten for this call - (array([1, 1, 1]), array([ 1, -1, 1, 1]), array([1, 1, 1, 1])) - >>> circuit(0.8) # back to default of 10 samples - array([ 1, -1, 1, 1, -1, 1, 1, 1, 1, 1]) - - When constructing a device, we may optionally pass a dictionary of custom - decompositions to be applied to certain operations upon device execution. - This is useful for enabling support of gates on devices where they would normally - be unsupported. - - For example, suppose we are running on an ion trap device which does not - natively implement the CNOT gate, but we would still like to write our - circuits in terms of CNOTs. On a ion trap device, CNOT can be implemented - using the ``IsingXX`` gate. We first define a decomposition function - (such functions have the signature ``decomposition(*params, wires)``): - - .. code-block:: python - - def ion_trap_cnot(wires, **_): - return [ - qml.RY(np.pi/2, wires=wires[0]), - qml.IsingXX(np.pi/2, wires=wires), - qml.RX(-np.pi/2, wires=wires[0]), - qml.RY(-np.pi/2, wires=wires[0]), - qml.RY(-np.pi/2, wires=wires[1]) - ] - - Next, we create a device, and a QNode for testing. When constructing the - QNode, we can set the expansion strategy to ``"device"`` to ensure the - decomposition is applied and will be viewable when we draw the circuit. - Note that custom decompositions should accept keyword arguments even when - it is not used. - - .. code-block:: python - - # As the CNOT gate normally has no decomposition, we can use default.qubit - # here for expository purposes. - dev = qml.device( - 'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot} - ) - - @qml.qnode(dev, expansion_strategy="device") - def run_cnot(): - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.X(1)) - - >>> print(qml.draw(run_cnot)()) - 0: ──RY(1.57)─╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)── - 1: ───────────╰IsingXX(1.57)──RY(-1.57)───────────── - - Some devices may accept additional arguments. For instance, - ``default.gaussian`` accepts the keyword argument ``hbar``, to set - the convention used in the commutation relation :math:`[\x,\p]=i\hbar` - (by default set to 2). - - Please refer to the documentation for the individual devices to see any - additional arguments that might be required or supported. - """ - if name not in plugin_devices: - # Device does not exist in the loaded device list. - # Attempt to refresh the devices, in case the user - # installed the plugin during the current Python session. - refresh_devices() - - if name in plugin_devices: - options = {} - - # load global configuration settings if available - config = kwargs.get("config", default_config) - - if config: - # combine configuration options with keyword arguments. - # Keyword arguments take preference, followed by device options, - # followed by plugin options, followed by global options. - options.update(config["main"]) - options.update(config[name.split(".")[0] + ".global"]) - options.update(config[name]) - - # Pop the custom decomposition keyword argument; we will use it here - # only and not pass it to the device. - custom_decomps = kwargs.pop("custom_decomps", None) - decomp_depth = kwargs.pop("decomp_depth", 10) - - kwargs.pop("config", None) - options.update(kwargs) - - # loads the device class - plugin_device_class = plugin_devices[name].load() - - if hasattr(plugin_device_class, "pennylane_requires") and Version( - version() - ) not in SimpleSpec(plugin_device_class.pennylane_requires): - raise DeviceError( - f"The {name} plugin requires PennyLane versions {plugin_device_class.pennylane_requires}, " - f"however PennyLane version {__version__} is installed." - ) - - # Construct the device - dev = plugin_device_class(*args, **options) - - # Once the device is constructed, we set its custom expansion function if - # any custom decompositions were specified. - if custom_decomps is not None: - if isinstance(dev, pennylane.devices.LegacyDevice): - custom_decomp_expand_fn = pennylane.transforms.create_decomp_expand_fn( - custom_decomps, dev, decomp_depth=decomp_depth - ) - dev.custom_expand(custom_decomp_expand_fn) - else: - custom_decomp_preprocess = ( - pennylane.transforms.tape_expand._create_decomp_preprocessing( - custom_decomps, dev, decomp_depth=decomp_depth - ) - ) - dev.preprocess = custom_decomp_preprocess - - return dev - - raise DeviceError(f"Device {name} does not exist. Make sure the required plugin is installed.") + raise AttributeError(f"module 'pennylane' has no attribute '{name}'") def version(): diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py index 5123d407ba5..9b0ea67d890 100644 --- a/pennylane/_qubit_device.py +++ b/pennylane/_qubit_device.py @@ -26,7 +26,7 @@ import logging import warnings from collections import defaultdict -from typing import List, Union +from typing import Union import numpy as np @@ -283,10 +283,15 @@ def execute(self, circuit, **kwargs): shots=[1], trainable_params=circuit.trainable_params, ) + # Some devices like Lightning-Kokkos use `self.shots` to update `_samples`, + # and hence we update `self.shots` temporarily for this loop + shots_copy = self.shots + self.shots = 1 for _ in circuit.shots: kwargs["mid_measurements"] = {} self.reset() results.append(self.execute(aux_circ, **kwargs)) + self.shots = shots_copy return tuple(results) # apply all circuit operations self.apply( @@ -1772,7 +1777,7 @@ def _adjoint_jacobian_processing(jac): # must be 2-dimensional return tuple(tuple(np.array(j_) for j_ in j) for j in jac) - def _get_diagonalizing_gates(self, circuit: QuantumTape) -> List[Operation]: + def _get_diagonalizing_gates(self, circuit: QuantumTape) -> list[Operation]: """Returns the gates that diagonalize the measured wires such that they are in the eigenbasis of the circuit observables. diff --git a/pennylane/capture/capture_qnode.py b/pennylane/capture/capture_qnode.py index d0faf65b665..0b7ecb6764a 100644 --- a/pennylane/capture/capture_qnode.py +++ b/pennylane/capture/capture_qnode.py @@ -15,6 +15,7 @@ This submodule defines a capture compatible call to QNodes. """ +import warnings from copy import copy from dataclasses import asdict from functools import lru_cache, partial @@ -29,7 +30,7 @@ def _get_shapes_for(*measurements, shots=None, num_device_wires=0): - if jax.config.jax_enable_x64: + if jax.config.jax_enable_x64: # pylint: disable=no-member dtype_map = { float: jax.numpy.float64, int: jax.numpy.int64, @@ -65,7 +66,13 @@ def _(*args, qnode, shots, device, qnode_kwargs, qfunc_jaxpr): def qfunc(*inner_args): return jax.core.eval_jaxpr(qfunc_jaxpr.jaxpr, qfunc_jaxpr.consts, *inner_args) - qnode = qml.QNode(qfunc, device, **qnode_kwargs) + with warnings.catch_warnings(): + warnings.filterwarnings( + action="ignore", + message=r"The max_expansion argument is deprecated and will be removed in version 0.39.", + category=qml.PennyLaneDeprecationWarning, + ) + qnode = qml.QNode(qfunc, device, **qnode_kwargs) return qnode._impl_call(*args, shots=shots) # pylint: disable=protected-access # pylint: disable=unused-argument diff --git a/pennylane/capture/primitives.py b/pennylane/capture/primitives.py index c8affc0a445..a26f3aac8e3 100644 --- a/pennylane/capture/primitives.py +++ b/pennylane/capture/primitives.py @@ -15,8 +15,9 @@ This submodule defines the abstract classes and primitives for capture. """ +from collections.abc import Callable from functools import lru_cache -from typing import Callable, Optional, Tuple, Type +from typing import Optional, Type import pennylane as qml @@ -107,7 +108,7 @@ def __init__( self._n_wires = n_wires self.has_eigvals: bool = has_eigvals - def abstract_eval(self, num_device_wires: int, shots: int) -> Tuple[Tuple, type]: + def abstract_eval(self, num_device_wires: int, shots: int) -> tuple[tuple, type]: """Calculate the shape and dtype for an evaluation with specified number of device wires and shots. diff --git a/pennylane/capture/switches.py b/pennylane/capture/switches.py index cbb25855c87..12e998dcb6c 100644 --- a/pennylane/capture/switches.py +++ b/pennylane/capture/switches.py @@ -15,7 +15,7 @@ Contains the switches to (de)activate the capturing mechanism, and a status reporting function on whether it is enabled or not. """ -from typing import Callable +from collections.abc import Callable has_jax = True try: @@ -24,7 +24,7 @@ has_jax = False -def _make_switches() -> [Callable[[], None], Callable[[], None], Callable[[], bool]]: +def _make_switches() -> tuple[Callable[[], None], Callable[[], None], Callable[[], bool]]: r"""Create three functions, corresponding to an activation switch, a deactivation switch and a status query, in that order. diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py index 69814d865ed..c93be7f74c4 100644 --- a/pennylane/circuit_graph.py +++ b/pennylane/circuit_graph.py @@ -15,42 +15,23 @@ This module contains the CircuitGraph class which is used to generate a DAG (directed acyclic graph) representation of a quantum circuit from an Operator queue. """ -from collections import namedtuple - -# pylint: disable=too-many-branches,too-many-arguments,too-many-instance-attributes -from numbers import Number +from collections import defaultdict, namedtuple +from functools import cached_property +from typing import List, Optional, Union import numpy as np import rustworkx as rx from pennylane.measurements import MeasurementProcess +from pennylane.operation import Observable, Operator +from pennylane.ops.identity import I +from pennylane.queuing import QueuingManager, WrappedObj from pennylane.resource import ResourcesOperation +from pennylane.wires import Wires -def _by_idx(x): - """Sorting key for Operators: queue index aka temporal order. - - Args: - x (Operator): node in the circuit graph - Returns: - int: sorting key for the node - """ - return x.queue_idx - - -def _is_observable(x): - """Predicate for deciding if an Operator instance is an observable. - - .. note:: - Currently some :class:`Observable` instances are not observables in this sense, - since they can be used as gates as well. - - Args: - x (Operator): node in the circuit graph - Returns: - bool: True iff x is an observable - """ - return isinstance(x, MeasurementProcess) +def _get_wires(obj, all_wires): + return all_wires if len(obj.wires) == 0 else obj.wires Layer = namedtuple("Layer", ["ops", "param_inds"]) @@ -74,6 +55,25 @@ def _is_observable(x): """ +def _construct_graph_from_queue(queue, all_wires): + inds_for_objs = defaultdict(list) # dict from wrappedobjs to all indices for the objs + nodes_on_wires = defaultdict(list) # wire to list of nodes + + graph = rx.PyDiGraph(multigraph=False) + + for i, obj in enumerate(queue): + inds_for_objs[WrappedObj(obj)].append(i) + + obj_node = graph.add_node(i) + for w in _get_wires(obj, all_wires): + if w in nodes_on_wires: + graph.add_edge(nodes_on_wires[w][-1], obj_node, "") + nodes_on_wires[w].append(obj_node) + + return graph, inds_for_objs, nodes_on_wires + + +# pylint: disable=too-many-instance-attributes, too-many-public-methods class CircuitGraph: """Represents a quantum circuit as a directed acyclic graph. @@ -84,85 +84,40 @@ class CircuitGraph: Args: ops (Iterable[.Operator]): quantum operators constituting the circuit, in temporal order - obs (Iterable[.MeasurementProcess]): terminal measurements, in temporal order + obs (List[Union[MeasurementProcess, Observable]]): terminal measurements, in temporal order wires (.Wires): The addressable wire registers of the device that will be executing this graph - par_info (list[dict]): Parameter information. For each index, the entry is a dictionary containing an operation + par_info (Optional[list[dict]]): Parameter information. For each index, the entry is a dictionary containing an operation and an index into that operation's parameters. - trainable_params (set[int]): A set containing the indices of parameters that support + trainable_params (Optional[set[int]]): A set containing the indices of parameters that support differentiability. The indices provided match the order of appearence in the quantum circuit. """ - # pylint: disable=too-many-public-methods - - def __init__(self, ops, obs, wires, par_info=None, trainable_params=None): + # pylint: disable=too-many-arguments + def __init__( + self, + ops: list[Union[Operator, MeasurementProcess]], + obs: List[Union[MeasurementProcess, Observable]], + wires: Wires, + par_info: Optional[list[dict]] = None, + trainable_params: Optional[set[int]] = None, + ): self._operations = ops self._observables = obs self.par_info = par_info self.trainable_params = trainable_params - queue = ops + obs + self._queue = ops + obs - self._depth = None - - self._grid = {} - """dict[int, list[Operator]]: dictionary representing the quantum circuit as a grid. - Here, the key is the wire number, and the value is a list containing the operators on that wire. - """ - - self._indices = {} - # Store indices for the nodes of the DAG here - - self.wires = wires + self.wires = Wires(wires) """Wires: wires that are addressed in the operations. Required to translate between wires and indices of the wires on the device.""" self.num_wires = len(wires) """int: number of wires the circuit contains""" - for k, op in enumerate(queue): - # meas_wires = wires or None # cannot use empty wire list in MeasurementProcess - op.queue_idx = k # store the queue index in the Operator - - for w in wires if len(op.wires) == 0 else op.wires: - # get the index of the wire on the device - wire = wires.index(w) - # add op to the grid, to the end of wire w - self._grid.setdefault(wire, []).append(op) - - # TODO: State preparations demolish the incoming state entirely, and therefore should have no incoming edges. - - self._graph = rx.PyDiGraph( - multigraph=False - ) #: rx.PyDiGraph: DAG representation of the quantum circuit - - # Iterate over each (populated) wire in the grid - for wire in self._grid.values(): - # Add the first operator on the wire to the graph - # This operator does not depend on any others - - # Check if wire[0] in self._grid.values() - # is already added to the graph; this - # condition avoids adding new nodes with - # the same value but different indexes - if all(wire[0] is not op for op in self._graph.nodes()): - _ind = self._graph.add_node(wire[0]) - self._indices.setdefault(id(wire[0]), _ind) - - for i in range(1, len(wire)): - # For subsequent operators on the wire: - if all(wire[i] is not op for op in self._graph.nodes()): - # Add them to the graph if they are not already - # in the graph (multi-qubit operators might already have been placed) - _ind = self._graph.add_node(wire[i]) - self._indices.setdefault(id(wire[i]), _ind) - - # Create an edge between this and the previous operator - # There isn't any default value for the edge-data in - # rx.PyDiGraph.add_edge(); this is set to an empty string - self._graph.add_edge(self._indices[id(wire[i - 1])], self._indices[id(wire[i])], "") - - # For computing depth; want only a graph with the operations, not - # including the observables - self._operation_graph = None + + self._graph, self._inds_for_objs, self._nodes_on_wires = _construct_graph_from_queue( + self._queue, wires + ) # Required to keep track if we need to handle multiple returned # observables per wire @@ -181,7 +136,7 @@ def print_contents(self): for op in self.observables: print(repr(op)) - def serialize(self): + def serialize(self) -> str: """Serialize the quantum circuit graph based on the operations and observables in the circuit graph and the index of the variables used by them. @@ -223,7 +178,7 @@ def serialize(self): return serialization_string @property - def hash(self): + def hash(self) -> int: """Creating a hash for the circuit graph based on the string generated by serialize. Returns: @@ -240,10 +195,9 @@ def observables_in_order(self): Currently the topological order is determined by the queue index. Returns: - list[Observable]: observables + List[Union[MeasurementProcess, Observable]]: observables """ - nodes = [node for node in self._graph.nodes() if _is_observable(node)] - return sorted(nodes, key=_by_idx) + return self._observables @property def observables(self): @@ -262,8 +216,7 @@ def operations_in_order(self): Returns: list[Operation]: operations """ - nodes = [node for node in self._graph.nodes() if not _is_observable(node)] - return sorted(nodes, key=_by_idx) + return self._operations @property def operations(self): @@ -274,7 +227,7 @@ def operations(self): def graph(self): """The graph representation of the quantum circuit. - The graph has nodes representing :class:`.Operator` instances, + The graph has nodes representing indices into the queue, and directed edges pointing from nodes to their immediate dependents/successors. Returns: @@ -291,9 +244,9 @@ def wire_indices(self, wire): Returns: list[int]: indices of operators on the wire, in temporal order """ - return [op.queue_idx for op in self._grid[wire]] + return self._nodes_on_wires[wire] - def ancestors(self, ops): + def ancestors(self, ops, sort=False): """Ancestors of a given set of operators. Args: @@ -302,14 +255,25 @@ def ancestors(self, ops): Returns: list[Operator]: ancestors of the given operators """ - # rx.ancestors() returns node indices instead of node-values - all_indices = set().union(*(rx.ancestors(self._graph, self._indices[id(o)]) for o in ops)) - double_op_indices = set(self._indices[id(o)] for o in ops) - ancestor_indices = all_indices - double_op_indices - - return list(self._graph.get_node_data(n) for n in ancestor_indices) - - def descendants(self, ops): + if isinstance(ops, (Operator, MeasurementProcess)): + raise ValueError( + "CircuitGraph.ancestors accepts an iterable of" + " operators and measurements, not operators and measurements themselves." + ) + if any(len(self._inds_for_objs[WrappedObj(op)]) > 1 for op in ops): + raise ValueError( + "Cannot calculate ancestors for an operator that occurs multiple times." + ) + ancestors = set() + for op in ops: + ind = self._inds_for_objs[WrappedObj(op)][0] + op_ancestors = rx.ancestors(self._graph, ind) + ancestors.update(set(op_ancestors)) + if sort: + ancestors = sorted(ancestors) + return [self._queue[ind] for ind in ancestors] + + def descendants(self, ops, sort=False): """Descendants of a given set of operators. Args: @@ -318,25 +282,23 @@ def descendants(self, ops): Returns: list[Operator]: descendants of the given operators """ - # rx.descendants() returns node indices instead of node-values - all_indices = set().union(*(rx.descendants(self._graph, self._indices[id(o)]) for o in ops)) - double_op_indices = set(self._indices[id(o)] for o in ops) - ancestor_indices = all_indices - double_op_indices - - return list(self._graph.get_node_data(n) for n in ancestor_indices) - - def _in_topological_order(self, ops): - """Sorts a set of operators in the circuit in a topological order. - - Args: - ops (Iterable[Operator]): set of operators in the circuit - - Returns: - Iterable[Operator]: same set of operators, topologically ordered - """ - G = self._graph.subgraph(list(self._indices[id(o)] for o in ops)) - indexes = rx.topological_sort(G) - return list(G[x] for x in indexes) + if isinstance(ops, (Operator, MeasurementProcess)): + raise ValueError( + "CircuitGraph.descendants accepts an iterable of" + " operators and measurements, not operators and measurements themselves." + ) + if any(len(self._inds_for_objs[WrappedObj(op)]) > 1 for op in ops): + raise ValueError( + "cannot calculate decendents for an operator that occurs multiple times." + ) + descendants = set() + for op in ops: + ind = self._inds_for_objs[WrappedObj(op)][0] + op_descendants = rx.descendants(self._graph, ind) + descendants.update(set(op_descendants)) + if sort: + descendants = sorted(descendants) + return [self._queue[ind] for ind in descendants] def ancestors_in_order(self, ops): """Operator ancestors in a topological order. @@ -349,7 +311,7 @@ def ancestors_in_order(self, ops): Returns: list[Operator]: ancestors of the given operators, topologically ordered """ - return sorted(self.ancestors(ops), key=_by_idx) # an abitrary topological order + return self.ancestors(ops, sort=True) def descendants_in_order(self, ops): """Operator descendants in a topological order. @@ -362,7 +324,7 @@ def descendants_in_order(self, ops): Returns: list[Operator]: descendants of the given operators, topologically ordered """ - return sorted(self.descendants(ops), key=_by_idx) + return self.descendants(ops, sort=True) def nodes_between(self, a, b): r"""Nodes on all the directed paths between the two given nodes. @@ -439,84 +401,74 @@ def update_node(self, old, new): Raises: ValueError: if the new :class:`~.Operator` does not act on the same wires as the old one """ - # NOTE Does not alter the graph edges in any way. variable_deps is not changed, _grid is not changed. Dangerous! + # NOTE Does not alter the graph edges in any way. variable_deps is not changed, Dangerous! if new.wires != old.wires: raise ValueError("The new Operator must act on the same wires as the old one.") - new.queue_idx = old.queue_idx - self._graph[self._indices[id(old)]] = new - index = self._indices.pop(id(old)) - self._indices[id(new)] = index + self._inds_for_objs[WrappedObj(new)] = self._inds_for_objs.pop(WrappedObj(old)) - self._operations = self.operations_in_order - self._observables = self.observables_in_order + for i, op in enumerate(self._operations): + if op is old: + self._operations[i] = new + for i, mp in enumerate(self._observables): + if mp is old: + self._observables[i] = new + for i, obj in enumerate(self._queue): + if obj is old: + self._queue[i] = new def get_depth(self): """Depth of the quantum circuit (longest path in the DAG).""" + return self._depth + + @cached_property + def _depth(self): # If there are no operations in the circuit, the depth is 0 if not self.operations: - self._depth = 0 - - # If there are operations but depth is uncomputed, compute the truncated graph - # with only the operations, and return the longest path + 1 (since the path is - # expressed in terms of edges, and we want it in terms of nodes). - if self._depth is None and self.operations: - if self._operation_graph is None: - self._operation_graph = self._graph.subgraph( - list(self._indices[id(node)] for node in self.operations) - ) - self._extend_graph(self._operation_graph) - self._depth = ( - rx.dag_longest_path_length( - self._operation_graph, weight_fn=lambda _, __, w: self._weight_func(w) - ) - + 1 - ) - return self._depth + return 0 + with QueuingManager.stop_recording(): + ops_with_initial_I = [ + I(self.wires) + ] + self.operations # add identity wire to end the graph + operation_graph, _, _ = _construct_graph_from_queue(ops_with_initial_I, self.wires) + + # pylint: disable=unused-argument + def weight_fn(in_idx, out_idx, w): + out_op = ops_with_initial_I[out_idx] + if isinstance(out_op, ResourcesOperation): + return out_op.resources().depth + return 1 + + return rx.dag_longest_path_length(operation_graph, weight_fn=weight_fn) + + def has_path_idx(self, a_idx: int, b_idx: int) -> bool: + """Checks if a path exists between the two given nodes. - @staticmethod - def _weight_func(weight): - """If weight is a number, use it!""" - if isinstance(weight, Number): - return weight - return 1 + Args: + a_idx (int): initial node index + b_idx (int): final node index - def _extend_graph(self, graph: rx.PyDiGraph) -> rx.PyDiGraph: - """Extend graph to account for custom depth operations""" - custom_depth_node_dict = {} - for op in self.operations: - if isinstance(op, ResourcesOperation) and (d := op.resources().depth) > 1: - custom_depth_node_dict[graph.nodes().index(op)] = d - - def _link_graph(target_index, sub_graph, node_index): - """Link incoming and outgoing edges for the initial node to the sub-graph""" - if target_index == node_index: - return sub_graph.nodes().index(f"{node_index}.0") - return sub_graph.nodes().index(f"{node_index}.1") - - for node_index, depth in custom_depth_node_dict.items(): - # Construct sub_graph: - sub_graph = rx.PyDiGraph() - source_node, target_node = (f"{node_index}.0", f"{node_index}.1") - - sub_graph.add_node(source_node) - sub_graph.add_node(target_node) - - sub_graph.add_edge( - sub_graph.nodes().index(source_node), - sub_graph.nodes().index(target_node), - depth - 1, # set edge weight as depth - 1 - ) + Returns: + bool: returns ``True`` if a path exists + """ + if a_idx == b_idx: + return True - graph.substitute_node_with_subgraph( - node_index, - sub_graph, - lambda _, t, __: _link_graph( - t, sub_graph, node_index # pylint: disable=cell-var-from-loop - ), + return ( + len( + rx.digraph_dijkstra_shortest_paths( + self._graph, + a_idx, + b_idx, + weight_fn=None, + default_weight=1.0, + as_undirected=False, + ) ) + != 0 + ) - def has_path(self, a, b): + def has_path(self, a, b) -> bool: """Checks if a path exists between the two given nodes. Args: @@ -526,15 +478,22 @@ def has_path(self, a, b): Returns: bool: returns ``True`` if a path exists """ + if a is b: return True + if any(len(self._inds_for_objs[WrappedObj(o)]) > 1 for o in (a, b)): + raise ValueError( + "CircuitGraph.has_path does not work with operations that have been repeated. " + "Consider using has_path_idx instead." + ) + return ( len( rx.digraph_dijkstra_shortest_paths( self._graph, - self._indices[id(a)], - self._indices[id(b)], + self._inds_for_objs[WrappedObj(a)][0], + self._inds_for_objs[WrappedObj(b)][0], weight_fn=None, default_weight=1.0, as_undirected=False, @@ -543,7 +502,7 @@ def has_path(self, a, b): != 0 ) - @property + @cached_property def max_simultaneous_measurements(self): """Returns the maximum number of measurements on any wire in the circuit graph. @@ -570,15 +529,11 @@ def max_simultaneous_measurements(self): Returns: int: the maximum number of measurements """ - if self._max_simultaneous_measurements is None: - all_wires = [] - - for obs in self.observables: - all_wires.extend(obs.wires.tolist()) - - a = np.array(all_wires) - _, counts = np.unique(a, return_counts=True) - self._max_simultaneous_measurements = ( - counts.max() if counts.size != 0 else 1 - ) # qml.state() will result in an empty array - return self._max_simultaneous_measurements + all_wires = [] + + for obs in self.observables: + all_wires.extend(obs.wires.tolist()) + + a = np.array(all_wires) + _, counts = np.unique(a, return_counts=True) + return counts.max() if counts.size != 0 else 1 # qml.state() will result in an empty array diff --git a/pennylane/compiler/qjit_api.py b/pennylane/compiler/qjit_api.py index 140bca19b4b..8d26b37422d 100644 --- a/pennylane/compiler/qjit_api.py +++ b/pennylane/compiler/qjit_api.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """QJIT compatible quantum and compilation operations API""" +from collections.abc import Callable from .compiler import ( AvailableCompilers, @@ -377,20 +378,54 @@ def loop_rx(x): ops_loader = compilers[active_jit]["ops"].load() return ops_loader.while_loop(cond_fn) - raise CompileError("There is no active compiler package.") # pragma: no cover + # if there is no active compiler, simply interpret the while loop + # via the Python interpretor. + def _decorator(body_fn: Callable) -> Callable: + """Transform that will call the input ``body_fn`` until the closure variable ``cond_fn`` is met. + Args: + body_fn (Callable): -def for_loop(lower_bound, upper_bound, step): - """A :func:`~.qjit` compatible for-loop for PennyLane programs. + Closure Variables: + cond_fn (Callable): - .. note:: + Returns: + Callable: a callable with the same signature as ``body_fn`` and ``cond_fn``. + """ + return WhileLoopCallable(cond_fn, body_fn) - This function only supports the Catalyst compiler. See - :func:`catalyst.for_loop` for more details. + return _decorator - Please see the Catalyst :doc:`quickstart guide `, - as well as the :doc:`sharp bits and debugging tips ` - page for an overview of the differences between Catalyst and PennyLane. + +class WhileLoopCallable: # pylint:disable=too-few-public-methods + """Base class to represent a while loop. This class + when called with an initial state will execute the while + loop via the Python interpreter. + + Args: + cond_fn (Callable): the condition function in the while loop + body_fn (Callable): the function that is executed within the while loop + """ + + def __init__(self, cond_fn, body_fn): + self.cond_fn = cond_fn + self.body_fn = body_fn + + def __call__(self, *init_state): + args = init_state + fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None + + while self.cond_fn(*args): + fn_res = self.body_fn(*args) + args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else () + + return fn_res + + +def for_loop(lower_bound, upper_bound, step): + """A :func:`~.qjit` compatible for-loop for PennyLane programs. When + used without :func:`~.qjit`, this function will fall back to a standard + Python for loop. This decorator provides a functional version of the traditional for-loop, similar to `jax.cond.fori_loop `__. @@ -430,19 +465,14 @@ def for_loop(lower_bound, upper_bound, step, loop_fn, *args): across iterations is handled automatically by the provided loop bounds, it must not be returned from the function. - Raises: - CompileError: if the compiler is not installed - .. seealso:: :func:`~.while_loop`, :func:`~.qjit` **Example** - .. code-block:: python dev = qml.device("lightning.qubit", wires=1) - @qml.qjit @qml.qnode(dev) def circuit(n: int, x: float): @@ -457,10 +487,24 @@ def loop_rx(i, x): # apply the for loop final_x = loop_rx(x) - return qml.expval(qml.Z(0)), final_x + return qml.expval(qml.Z(0)) >>> circuit(7, 1.6) - (array(0.97926626), array(0.55395718)) + array(0.97926626) + + ``for_loop`` is also :func:`~.qjit` compatible; when used with the + :func:`~.qjit` decorator, the for loop will not be unrolled, and instead + will be captured as-is during compilation and executed during runtime: + + >>> qml.qjit(circuit)(7, 1.6) + Array(0.97926626, dtype=float64) + + .. note:: + + Please see the Catalyst :doc:`quickstart guide `, + as well as the :doc:`sharp bits and debugging tips ` + page for an overview of using quantum just-in-time compilation. + """ if active_jit := active_compiler(): @@ -468,4 +512,59 @@ def loop_rx(i, x): ops_loader = compilers[active_jit]["ops"].load() return ops_loader.for_loop(lower_bound, upper_bound, step) - raise CompileError("There is no active compiler package.") # pragma: no cover + # if there is no active compiler, simply interpret the for loop + # via the Python interpretor. + def _decorator(body_fn): + """Transform that will call the input ``body_fn`` within a for loop defined by the closure variables lower_bound, upper_bound, and step. + + Args: + body_fn (Callable): The function called within the for loop. Note that the loop body + function must always have the iteration index as its first + argument, which can be used arbitrarily inside the loop body. As the value of the index + across iterations is handled automatically by the provided loop bounds, it must not be + returned from the function. + + Closure Variables: + lower_bound (int): starting value of the iteration index + upper_bound (int): (exclusive) upper bound of the iteration index + step (int): increment applied to the iteration index at the end of each iteration + + Returns: + Callable: a callable with the same signature as ``body_fn`` + """ + return ForLoopCallable(lower_bound, upper_bound, step, body_fn) + + return _decorator + + +class ForLoopCallable: # pylint:disable=too-few-public-methods + """Base class to represent a for loop. This class + when called with an initial state will execute the while + loop via the Python interpreter. + + Args: + lower_bound (int): starting value of the iteration index + upper_bound (int): (exclusive) upper bound of the iteration index + step (int): increment applied to the iteration index at the end of each iteration + body_fn (Callable): The function called within the for loop. Note that the loop body + function must always have the iteration index as its first + argument, which can be used arbitrarily inside the loop body. As the value of the index + across iterations is handled automatically by the provided loop bounds, it must not be + returned from the function. + """ + + def __init__(self, lower_bound, upper_bound, step, body_fn): + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.step = step + self.body_fn = body_fn + + def __call__(self, *init_state): + args = init_state + fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None + + for i in range(self.lower_bound, self.upper_bound, self.step): + fn_res = self.body_fn(i, *args) + args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else () + + return fn_res diff --git a/pennylane/data/attributes/dictionary.py b/pennylane/data/attributes/dictionary.py index 443a566e90e..2fb45435fd1 100644 --- a/pennylane/data/attributes/dictionary.py +++ b/pennylane/data/attributes/dictionary.py @@ -15,9 +15,8 @@ of Dataset attributes.""" -import typing -from collections.abc import Mapping -from typing import Dict, Generic, Union +from collections.abc import Iterator, Mapping, MutableMapping +from typing import Generic, Union from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Any, HDF5Group @@ -27,8 +26,8 @@ class DatasetDict( # pylint: disable=too-many-ancestors Generic[T], - DatasetAttribute[HDF5Group, typing.Mapping[str, T], typing.Mapping[str, T]], - typing.MutableMapping[str, T], + DatasetAttribute[HDF5Group, Mapping[str, T], Mapping[str, T]], + MutableMapping[str, T], MapperMixin, ): """Provides a dict-like collection for Dataset attribute types. Keys must @@ -36,15 +35,15 @@ class DatasetDict( # pylint: disable=too-many-ancestors type_id = "dict" - def __post_init__(self, value: typing.Mapping[str, T]): + def __post_init__(self, value: Mapping[str, T]): super().__post_init__(value) self.update(value) @classmethod - def default_value(cls) -> Dict: + def default_value(cls) -> dict: return {} - def hdf5_to_value(self, bind: HDF5Group) -> typing.MutableMapping[str, T]: + def hdf5_to_value(self, bind: HDF5Group) -> MutableMapping[str, T]: return self def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: None) -> HDF5Group: @@ -52,10 +51,10 @@ def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: None) -> HDF5Gr return grp - def copy_value(self) -> Dict[str, T]: + def copy_value(self) -> dict[str, T]: return {key: attr.copy_value() for key, attr in self._mapper.items()} - def copy(self) -> Dict[str, T]: + def copy(self) -> dict[str, T]: """Returns a copy of this mapping as a builtin ``dict``, with all elements copied.""" return self.copy_value() @@ -93,7 +92,7 @@ def __eq__(self, __value: object) -> bool: return all(__value[key] == self[key] for key in __value.keys()) - def __iter__(self) -> typing.Iterator[str]: + def __iter__(self) -> Iterator[str]: return (key for key in self.bind.keys()) def __str__(self) -> str: diff --git a/pennylane/data/attributes/list.py b/pennylane/data/attributes/list.py index 098bddd5eef..b721f76a740 100644 --- a/pennylane/data/attributes/list.py +++ b/pennylane/data/attributes/list.py @@ -14,9 +14,8 @@ """Contains an DatasetAttribute that allows for heterogeneous lists of dataset types.""" -import typing -from collections.abc import Sequence -from typing import Generic, List, Union, overload +from collections.abc import Iterable, MutableSequence, Sequence +from typing import Generic, Union, overload from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Any, HDF5Group @@ -26,37 +25,35 @@ class DatasetList( # pylint: disable=too-many-ancestors Generic[T], - DatasetAttribute[HDF5Group, typing.Sequence[T], typing.Iterable[T]], - typing.MutableSequence[T], + DatasetAttribute[HDF5Group, Sequence[T], Iterable[T]], + MutableSequence[T], MapperMixin, ): """Provides a list-like collection type for Dataset Attributes.""" type_id = "list" - def __post_init__(self, value: typing.Iterable[T]): + def __post_init__(self, value: Iterable[T]): super().__post_init__(value) self.extend(value) @classmethod - def default_value(cls) -> typing.Iterable[T]: + def default_value(cls) -> Iterable[T]: return [] - def hdf5_to_value(self, bind: HDF5Group) -> typing.MutableSequence[T]: + def hdf5_to_value(self, bind: HDF5Group) -> MutableSequence[T]: return self - def value_to_hdf5( - self, bind_parent: HDF5Group, key: str, value: typing.Iterable[T] - ) -> HDF5Group: + def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: Iterable[T]) -> HDF5Group: grp = bind_parent.create_group(key) return grp - def copy_value(self) -> List[T]: + def copy_value(self) -> list[T]: return [self._mapper[str(i)].copy_value() for i in range(len(self))] - def copy(self) -> List[T]: + def copy(self) -> list[T]: """Returns a copy of this list as a builtin ``list``, with all elements copied..""" return self.copy_value() @@ -97,7 +94,7 @@ def __repr__(self) -> str: return f"[{items_repr}]" @overload - def __getitem__(self, index: slice) -> typing.List[T]: + def __getitem__(self, index: slice) -> list[T]: pass @overload diff --git a/pennylane/data/attributes/molecule.py b/pennylane/data/attributes/molecule.py index ed09c75f174..362da360a84 100644 --- a/pennylane/data/attributes/molecule.py +++ b/pennylane/data/attributes/molecule.py @@ -13,7 +13,7 @@ # limitations under the License. """Contains DatasetAttribute definition for ``pennylane.qchem.Molecule``.""" -from typing import Tuple, Type +from typing import Type from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Group @@ -27,7 +27,7 @@ class DatasetMolecule(DatasetAttribute[HDF5Group, Molecule, Molecule]): type_id = "molecule" @classmethod - def consumes_types(cls) -> Tuple[Type[Molecule]]: + def consumes_types(cls) -> tuple[Type[Molecule]]: return (Molecule,) def hdf5_to_value(self, bind: HDF5Group) -> Molecule: diff --git a/pennylane/data/attributes/none.py b/pennylane/data/attributes/none.py index cd13c4fe92d..ab234acd287 100644 --- a/pennylane/data/attributes/none.py +++ b/pennylane/data/attributes/none.py @@ -13,7 +13,7 @@ # limitations under the License. """Contains DatasetAttribute definition for None""" -from typing import Literal, Tuple, Type +from typing import Literal, Type from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Array, HDF5Group @@ -29,7 +29,7 @@ def default_value(cls) -> Literal[None]: return None @classmethod - def consumes_types(cls) -> Tuple[Type[None]]: + def consumes_types(cls) -> tuple[Type[None]]: return (type(None),) def hdf5_to_value(self, bind) -> None: diff --git a/pennylane/data/attributes/operator/operator.py b/pennylane/data/attributes/operator/operator.py index 28604b2cfc1..5301654ab9d 100644 --- a/pennylane/data/attributes/operator/operator.py +++ b/pennylane/data/attributes/operator/operator.py @@ -15,9 +15,9 @@ of operators.""" import json -import typing +from collections.abc import Sequence from functools import lru_cache -from typing import Dict, FrozenSet, Generic, List, Type, TypeVar +from typing import Generic, Type, TypeVar import numpy as np @@ -52,7 +52,7 @@ class DatasetOperator(Generic[Op], DatasetAttribute[HDF5Group, Op, Op]): @classmethod @lru_cache(1) - def supported_ops(cls) -> FrozenSet[Type[Operator]]: + def supported_ops(cls) -> frozenset[Type[Operator]]: """Set of supported operators.""" return frozenset( ( @@ -198,7 +198,7 @@ def hdf5_to_value(self, bind: HDF5Group) -> Op: return self._hdf5_to_ops(bind)[0] def _ops_to_hdf5( - self, bind_parent: HDF5Group, key: str, value: typing.Sequence[Operator] + self, bind_parent: HDF5Group, key: str, value: Sequence[Operator] ) -> HDF5Group: """Serialize op sequence ``value``, and create nested sequences for any composite ops in ``value``. @@ -247,7 +247,7 @@ def _ops_to_hdf5( return bind - def _hdf5_to_ops(self, bind: HDF5Group) -> List[Operator]: + def _hdf5_to_ops(self, bind: HDF5Group) -> list[Operator]: """Load list of serialized ops from ``bind``.""" ops = [] @@ -293,6 +293,6 @@ def _hdf5_to_ops(self, bind: HDF5Group) -> List[Operator]: @classmethod @lru_cache(1) - def _supported_ops_dict(cls) -> Dict[str, Type[Operator]]: + def _supported_ops_dict(cls) -> dict[str, Type[Operator]]: """Returns a dict mapping ``Operator`` subclass names to the class.""" return {op.__name__: op for op in cls.supported_ops()} diff --git a/pennylane/data/attributes/sparse_array.py b/pennylane/data/attributes/sparse_array.py index 246f94e06bd..94b062d2f34 100644 --- a/pennylane/data/attributes/sparse_array.py +++ b/pennylane/data/attributes/sparse_array.py @@ -14,7 +14,7 @@ """Contains DatasetAttribute definition for ``scipy.sparse.csr_array``.""" from functools import lru_cache -from typing import Dict, Generic, Tuple, Type, TypeVar, Union, cast +from typing import Generic, Type, TypeVar, Union, cast import numpy as np from scipy.sparse import ( @@ -64,7 +64,7 @@ def sparse_array_class(self) -> Type[SparseT]: @classmethod def consumes_types( cls, - ) -> Tuple[Type[Union[SparseArray, SparseMatrix]], ...]: + ) -> tuple[Type[Union[SparseArray, SparseMatrix]], ...]: return ( bsr_array, coo_array, @@ -122,6 +122,6 @@ def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: SparseT) -> HDF @classmethod @lru_cache(1) - def _supported_sparse_dict(cls) -> Dict[str, Type[Union[SparseArray, SparseMatrix]]]: + def _supported_sparse_dict(cls) -> dict[str, Type[Union[SparseArray, SparseMatrix]]]: """Returns a dict mapping sparse array class names to the class.""" return {op.__name__: op for op in cls.consumes_types()} diff --git a/pennylane/data/attributes/string.py b/pennylane/data/attributes/string.py index 87cec14f0cb..9a2155e7def 100644 --- a/pennylane/data/attributes/string.py +++ b/pennylane/data/attributes/string.py @@ -14,7 +14,7 @@ """Contains an DatasetAttribute for str objects.""" -from typing import Tuple, Type +from typing import Type from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Array, HDF5Group @@ -26,7 +26,7 @@ class DatasetString(DatasetAttribute[HDF5Array, str, str]): type_id = "string" @classmethod - def consumes_types(cls) -> Tuple[Type[str]]: + def consumes_types(cls) -> tuple[Type[str]]: return (str,) def hdf5_to_value(self, bind: HDF5Array) -> str: diff --git a/pennylane/data/attributes/tuple.py b/pennylane/data/attributes/tuple.py index 571e6c380f2..4a37ec6f664 100644 --- a/pennylane/data/attributes/tuple.py +++ b/pennylane/data/attributes/tuple.py @@ -14,8 +14,7 @@ """Contains an DatasetAttribute that allows for heterogeneous tuples of dataset types.""" -import typing -from typing import Generic +from typing import Generic, Type from pennylane.data.base.attribute import DatasetAttribute from pennylane.data.base.hdf5 import HDF5Group @@ -25,21 +24,21 @@ class DatasetTuple( Generic[T], - DatasetAttribute[HDF5Group, typing.Tuple[T], typing.Tuple[T]], + DatasetAttribute[HDF5Group, tuple[T], tuple[T]], ): """Type for tuples.""" type_id = "tuple" @classmethod - def consumes_types(cls) -> typing.Tuple[typing.Type[tuple]]: + def consumes_types(cls) -> tuple[Type[tuple]]: return (tuple,) @classmethod - def default_value(cls) -> typing.Tuple[()]: + def default_value(cls) -> tuple[()]: return tuple() - def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: typing.Tuple[T]) -> HDF5Group: + def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: tuple[T]) -> HDF5Group: grp = bind_parent.create_group(key) mapper = AttributeTypeMapper(grp) @@ -48,7 +47,7 @@ def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: typing.Tuple[T] return grp - def hdf5_to_value(self, bind: HDF5Group) -> typing.Tuple[T]: + def hdf5_to_value(self, bind: HDF5Group) -> tuple[T]: mapper = AttributeTypeMapper(bind) return tuple(mapper[str(i)].copy_value() for i in range(len(self.bind))) diff --git a/pennylane/data/base/_lazy_modules.py b/pennylane/data/base/_lazy_modules.py index 88db5b967d2..e2fe767be71 100644 --- a/pennylane/data/base/_lazy_modules.py +++ b/pennylane/data/base/_lazy_modules.py @@ -1,8 +1,9 @@ """Contains a lazy-loaded interface to the HDF5 module. For internal use only.""" import importlib +from collections.abc import Callable from types import ModuleType -from typing import Any, Callable, Optional, Union +from typing import Any, Optional, Union _MISSING_MODULES_EXC = ImportError( "This feature requires the 'aiohttp', 'h5py' and 'fsspec' packages. " diff --git a/pennylane/data/base/attribute.py b/pennylane/data/base/attribute.py index 2196ca5098f..e1eefbf9b07 100644 --- a/pennylane/data/base/attribute.py +++ b/pennylane/data/base/attribute.py @@ -14,26 +14,13 @@ """Contains the base class for Dataset attribute types, and a class for attribute metadata.""" -import typing import warnings from abc import ABC, abstractmethod -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence from functools import lru_cache from numbers import Number from types import MappingProxyType -from typing import ( - Any, - ClassVar, - Generic, - Iterator, - Literal, - Optional, - Tuple, - Type, - TypeVar, - Union, - overload, -) +from typing import Any, ClassVar, Generic, Literal, Optional, Type, TypeVar, Union, overload from pennylane.data.base import hdf5 from pennylane.data.base.hdf5 import HDF5, HDF5Any, HDF5Group @@ -55,12 +42,12 @@ class AttributeInfo(MutableMapping): """ attrs_namespace: ClassVar[str] = "qml.data" - attrs_bind: typing.MutableMapping[str, Any] + attrs_bind: MutableMapping[str, Any] @overload def __init__( # overload to specify known keyword args self, - attrs_bind: Optional[typing.MutableMapping[str, Any]] = None, + attrs_bind: Optional[MutableMapping[str, Any]] = None, *, doc: Optional[str] = None, py_type: Optional[str] = None, @@ -72,7 +59,7 @@ def __init__( # overload to specify known keyword args def __init__(self): # need at least two overloads when using @overload pass - def __init__(self, attrs_bind: Optional[typing.MutableMapping[str, Any]] = None, **kwargs: Any): + def __init__(self, attrs_bind: Optional[MutableMapping[str, Any]] = None, **kwargs: Any): object.__setattr__(self, "attrs_bind", attrs_bind if attrs_bind is not None else {}) for k, v in kwargs.items(): @@ -192,7 +179,7 @@ def __init__( value: Union[InitValueType, Literal[UNSET]] = UNSET, info: Optional[AttributeInfo] = None, *, - parent_and_key: Optional[Tuple[HDF5Group, str]] = None, + parent_and_key: Optional[tuple[HDF5Group, str]] = None, ): """Initialize a new dataset attribute from ``value``. @@ -220,7 +207,7 @@ def __init__( info: Optional[AttributeInfo] = None, *, bind: Optional[HDF5] = None, - parent_and_key: Optional[Tuple[HDF5Group, str]] = None, + parent_and_key: Optional[tuple[HDF5Group, str]] = None, ) -> None: """ Initialize a new dataset attribute, or load from an existing @@ -260,7 +247,7 @@ def _value_init( self, value: Union[InitValueType, Literal[UNSET]], info: Optional[AttributeInfo], - parent_and_key: Optional[Tuple[HDF5Group, str]], + parent_and_key: Optional[tuple[HDF5Group, str]], ): """Constructor for value initialization. See __init__().""" @@ -304,7 +291,7 @@ def py_type(cls, value_type: Type[InitValueType]) -> str: return get_type_str(value_type) @classmethod - def consumes_types(cls) -> typing.Iterable[type]: + def consumes_types(cls) -> Iterable[type]: """ Returns an iterable of types for which this should be the default codec. If a value of one of these types is assigned to a Dataset @@ -376,13 +363,13 @@ def __repr__(self) -> str: def __str__(self) -> str: return str(self.get_value()) - __registry: typing.Mapping[str, Type["DatasetAttribute"]] = {} - __type_consumer_registry: typing.Mapping[type, Type["DatasetAttribute"]] = {} + __registry: Mapping[str, Type["DatasetAttribute"]] = {} + __type_consumer_registry: Mapping[type, Type["DatasetAttribute"]] = {} - registry: typing.Mapping[str, Type["DatasetAttribute"]] = MappingProxyType(__registry) + registry: Mapping[str, Type["DatasetAttribute"]] = MappingProxyType(__registry) """Maps type_ids to their DatasetAttribute classes.""" - type_consumer_registry: typing.Mapping[type, Type["DatasetAttribute"]] = MappingProxyType( + type_consumer_registry: Mapping[type, Type["DatasetAttribute"]] = MappingProxyType( __type_consumer_registry ) """Maps types to their default DatasetAttribute""" diff --git a/pennylane/data/base/dataset.py b/pennylane/data/base/dataset.py index 89824fdf62d..6550716c375 100644 --- a/pennylane/data/base/dataset.py +++ b/pennylane/data/base/dataset.py @@ -16,24 +16,11 @@ for declaratively defining dataset classes. """ -import typing +from collections.abc import Iterable, Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import ( - Any, - ClassVar, - Generic, - List, - Literal, - Optional, - Tuple, - Type, - TypeVar, - Union, - cast, - get_origin, -) +from typing import Any, ClassVar, Generic, Literal, Optional, Type, TypeVar, Union, cast, get_origin # pylint doesn't think this exists from typing_extensions import dataclass_transform # pylint: disable=no-name-in-module @@ -154,9 +141,9 @@ class Dataset(MapperMixin, _DatasetTransform): """ __data_name__: ClassVar[str] - __identifiers__: ClassVar[Tuple[str, ...]] + __identifiers__: ClassVar[tuple[str, ...]] - fields: ClassVar[typing.Mapping[str, Field]] + fields: ClassVar[Mapping[str, Field]] """ A mapping of attribute names to their ``Attribute`` information. Note that this contains attributes declared on the class, not attributes added to @@ -171,7 +158,7 @@ def __init__( bind: Optional[HDF5Group] = None, *, data_name: Optional[str] = None, - identifiers: Optional[Tuple[str, ...]] = None, + identifiers: Optional[tuple[str, ...]] = None, **attrs: Any, ): """ @@ -245,7 +232,7 @@ def data_name(self) -> str: return self.info.get("data_name", self.__data_name__) @property - def identifiers(self) -> typing.Mapping[str, str]: # pylint: disable=function-redefined + def identifiers(self) -> Mapping[str, str]: # pylint: disable=function-redefined """Returns this dataset's parameters.""" return { attr_name: getattr(self, attr_name) @@ -264,12 +251,12 @@ def bind(self) -> HDF5Group: # pylint: disable=function-redefined return self._bind @property - def attrs(self) -> typing.Mapping[str, DatasetAttribute]: + def attrs(self) -> Mapping[str, DatasetAttribute]: """Returns all attributes of this Dataset.""" return self._mapper.view() @property - def attr_info(self) -> typing.Mapping[str, AttributeInfo]: + def attr_info(self) -> Mapping[str, AttributeInfo]: """Returns a mapping of the ``AttributeInfo`` for each of this dataset's attributes.""" return MappingProxyType( { @@ -278,14 +265,14 @@ def attr_info(self) -> typing.Mapping[str, AttributeInfo]: } ) - def list_attributes(self) -> List[str]: + def list_attributes(self) -> list[str]: """Returns a list of this dataset's attributes.""" return list(self.attrs.keys()) def read( self, source: Union[str, Path, "Dataset"], - attributes: Optional[typing.Iterable[str]] = None, + attributes: Optional[Iterable[str]] = None, *, overwrite: bool = False, ) -> None: @@ -311,7 +298,7 @@ def write( self, dest: Union[str, Path, "Dataset"], mode: Literal["w", "w-", "a"] = "a", - attributes: Optional[typing.Iterable[str]] = None, + attributes: Optional[Iterable[str]] = None, *, overwrite: bool = False, ) -> None: @@ -345,7 +332,7 @@ def write( hdf5.copy_all(self.bind, dest.bind, *missing_identifiers) def _init_bind( - self, data_name: Optional[str] = None, identifiers: Optional[Tuple[str, ...]] = None + self, data_name: Optional[str] = None, identifiers: Optional[tuple[str, ...]] = None ): if self.bind.file.mode == "r+": if "type_id" not in self.info: @@ -400,7 +387,7 @@ def __repr__(self) -> str: return f"<{type(self).__name__} = {repr_items}>" def __init_subclass__( - cls, *, data_name: Optional[str] = None, identifiers: Optional[Tuple[str, ...]] = None + cls, *, data_name: Optional[str] = None, identifiers: Optional[tuple[str, ...]] = None ) -> None: """Initializes the ``fields`` dict of a Dataset subclass using the declared ``Attributes`` and their type annotations.""" diff --git a/pennylane/data/base/mapper.py b/pennylane/data/base/mapper.py index ed431f26ae7..ac0e2808db6 100644 --- a/pennylane/data/base/mapper.py +++ b/pennylane/data/base/mapper.py @@ -15,10 +15,9 @@ class that provides the mapper class.""" -import typing -from collections.abc import MutableMapping +from collections.abc import Iterator, Mapping, MutableMapping from types import MappingProxyType -from typing import Any, Dict, Optional, Type +from typing import Any, Optional, Type from pennylane.data.base.attribute import ( AttributeInfo, @@ -46,7 +45,7 @@ class AttributeTypeMapper(MutableMapping): """ bind: HDF5Group - _cache: Dict[str, DatasetAttribute] + _cache: dict[str, DatasetAttribute] def __init__(self, bind: HDF5Group) -> None: self._cache = {} @@ -120,14 +119,14 @@ def move(self, src: str, dest: str) -> None: self.bind.move(src, dest) self._cache.pop(src, None) - def view(self) -> typing.Mapping[str, DatasetAttribute]: + def view(self) -> Mapping[str, DatasetAttribute]: """Returns a read-only mapping of the attributes in ``bind``.""" return MappingProxyType(self) def __len__(self) -> int: return len(self.bind) - def __iter__(self) -> typing.Iterator[str]: + def __iter__(self) -> Iterator[str]: return iter(self.bind) def __contains__(self, key: str) -> bool: diff --git a/pennylane/data/base/typing_util.py b/pennylane/data/base/typing_util.py index 18b3b1e1ee4..77a04a1dd26 100644 --- a/pennylane/data/base/typing_util.py +++ b/pennylane/data/base/typing_util.py @@ -18,12 +18,9 @@ from functools import lru_cache from typing import ( Any, - Dict, ForwardRef, - List, Literal, Optional, - Tuple, Type, TypeVar, Union, @@ -32,7 +29,7 @@ get_origin, ) -JSON = Union[str, int, bool, float, None, Dict[str, Any], List[Any]] +JSON = Union[str, int, bool, float, None, dict[str, Any], list[Any]] # Generic type variable T = TypeVar("T") @@ -123,7 +120,7 @@ def get_type_str(cls: Union[type, str, None]) -> str: # pylint: disable=too-man return f"{cls.__module__}.{cls.__qualname__}" -def resolve_special_type(type_: Any) -> Optional[Tuple[type, List[type]]]: +def resolve_special_type(type_: Any) -> Optional[tuple[type, list[type]]]: """Converts special typing forms (Union[...], Optional[...]), and parametrized generics (List[...], Dict[...]) into a 2-tuple of its base type and arguments. If ``type_`` is a regular type, or an object, this function will return diff --git a/pennylane/data/data_manager/__init__.py b/pennylane/data/data_manager/__init__.py index cfde2ed5e9b..b07e3017f94 100644 --- a/pennylane/data/data_manager/__init__.py +++ b/pennylane/data/data_manager/__init__.py @@ -16,13 +16,13 @@ them. """ -import typing import urllib.parse +from collections.abc import Mapping, Iterable from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from functools import lru_cache from pathlib import Path from time import sleep -from typing import List, Optional, Union +from typing import Optional, Union from requests import get @@ -58,7 +58,7 @@ def _get_data_struct(): def _download_partial( s3_url: str, dest: Path, - attributes: Optional[typing.Iterable[str]], + attributes: Optional[Iterable[str]], overwrite: bool, block_size: int, ) -> None: @@ -112,7 +112,7 @@ def _download_full(s3_url: str, dest: Path): def _download_dataset( data_path: DataPath, dest: Path, - attributes: Optional[typing.Iterable[str]], + attributes: Optional[Iterable[str]], block_size: int, force: bool = False, ) -> None: @@ -136,7 +136,7 @@ def _download_dataset( _download_full(s3_url, dest=dest) -def _validate_attributes(data_struct: dict, data_name: str, attributes: typing.Iterable[str]): +def _validate_attributes(data_struct: dict, data_name: str, attributes: Iterable[str]): """Checks that ``attributes`` contains only valid attributes for the given ``data_name``. If any attributes do not exist, raise a ValueError.""" invalid_attributes = [ @@ -155,12 +155,12 @@ def _validate_attributes(data_struct: dict, data_name: str, attributes: typing.I def load( # pylint: disable=too-many-arguments data_name: str, - attributes: Optional[typing.Iterable[str]] = None, + attributes: Optional[Iterable[str]] = None, folder_path: Path = Path("./datasets/"), force: bool = False, num_threads: int = 50, block_size: int = 8388608, - **params: Union[ParamArg, str, List[str]], + **params: Union[ParamArg, str, list[str]], ): r"""Downloads the data if it is not already present in the directory and returns it as a list of :class:`~pennylane.data.Dataset` objects. For the full list of available datasets, please see @@ -314,7 +314,7 @@ def remove_paths(foldermap): to Paths to a list of the parameters.""" value = next(iter(foldermap.values())) - if not isinstance(value, typing.Mapping): + if not isinstance(value, Mapping): return sorted(foldermap.keys()) return {param: remove_paths(foldermap[param]) for param in foldermap.keys()} diff --git a/pennylane/data/data_manager/foldermap.py b/pennylane/data/data_manager/foldermap.py index bf4c0cde4e0..48eba6e979f 100644 --- a/pennylane/data/data_manager/foldermap.py +++ b/pennylane/data/data_manager/foldermap.py @@ -17,10 +17,9 @@ """ -import typing -from collections.abc import Mapping +from collections.abc import Iterable, Iterator, Mapping from pathlib import PurePosixPath -from typing import Any, List, Literal, Optional, Tuple, Union +from typing import Any, Literal, Optional, Union from .params import Description, ParamArg, ParamVal @@ -33,7 +32,7 @@ def __repr__(self) -> str: return repr(str(self)) -class FolderMapView(typing.Mapping[str, Union["FolderMapView", DataPath]]): +class FolderMapView(Mapping[str, Union["FolderMapView", DataPath]]): """Provides a read-only view of the ``foldermap.json`` file in the datasets bucket. The folder map is a nested mapping of dataset parameters to their path, relative to the ``foldermap.json`` @@ -75,7 +74,7 @@ class FolderMapView(typing.Mapping[str, Union["FolderMapView", DataPath]]): __PRIVATE_KEYS = {"__default", "__params"} - def __init__(self, __curr_level: typing.Mapping[str, Any]) -> None: + def __init__(self, __curr_level: Mapping[str, Any]) -> None: """Initialize the mapping. Args: @@ -93,8 +92,8 @@ def find( self, data_name: str, missing_default: Optional[ParamArg] = ParamArg.DEFAULT, - **params: Union[typing.Iterable[ParamVal], ParamArg], - ) -> List[Tuple[Description, DataPath]]: + **params: Union[Iterable[ParamVal], ParamArg], + ) -> list[tuple[Description, DataPath]]: """Returns a 2-tuple of dataset description and paths, for each dataset that matches ``params``.""" @@ -104,15 +103,15 @@ def find( raise RuntimeError("Can only call find() from top level of foldermap") from exc try: - param_names: List[str] = data_names_to_params[data_name] + param_names: list[str] = data_names_to_params[data_name] except KeyError as exc: raise ValueError(f"No datasets with data name: '{data_name}'") from exc - curr: List[Tuple[Description, Union[FolderMapView, DataPath]]] = [ + curr: list[tuple[Description, Union[FolderMapView, DataPath]]] = [ (Description(()), self[data_name]) ] - todo: List[Tuple[Description, Union[FolderMapView, DataPath]]] = [] - done: List[Tuple[Description, DataPath]] = [] + todo: list[tuple[Description, Union[FolderMapView, DataPath]]] = [] + done: list[tuple[Description, DataPath]] = [] for param_name in param_names: param_arg = params.get(param_name, missing_default) @@ -183,10 +182,11 @@ def __getitem__( return DataPath(elem) - def __iter__(self) -> typing.Iterator[str]: + def __iter__(self) -> Iterator[str]: return (key for key in self.__curr_level.keys() if key not in self.__PRIVATE_KEYS) - def keys(self) -> typing.FrozenSet[str]: + def keys(self) -> frozenset[str]: + """Keys of the folder view""" return frozenset(iter(self)) def __len__(self) -> int: diff --git a/pennylane/data/data_manager/params.py b/pennylane/data/data_manager/params.py index f8f72b5fc64..0e467777cfd 100644 --- a/pennylane/data/data_manager/params.py +++ b/pennylane/data/data_manager/params.py @@ -17,9 +17,9 @@ import enum -import typing +from collections.abc import Iterable, Iterator, Mapping from functools import lru_cache -from typing import Any, Dict, FrozenSet, Iterable, List, Tuple, Union +from typing import Any, Union class ParamArg(enum.Enum): @@ -34,7 +34,7 @@ class ParamArg(enum.Enum): @classmethod @lru_cache(maxsize=1) - def values(cls) -> FrozenSet[str]: + def values(cls) -> frozenset[str]: """Returns all values.""" return frozenset(arg.value for arg in cls) @@ -57,18 +57,18 @@ def __str__(self) -> str: # pylint: disable=invalid-str-returned ParamVal = str -class Description(typing.Mapping[ParamName, ParamVal]): +class Description(Mapping[ParamName, ParamVal]): """An immutable and hashable dictionary that contains all the parameter values for a dataset.""" - def __init__(self, params: typing.Iterable[Tuple[ParamName, ParamVal]]): + def __init__(self, params: Iterable[tuple[ParamName, ParamVal]]): self.__data = dict(params) self.__hash = None def __getitem__(self, __key: ParamName) -> ParamVal: return self.__data[__key] - def __iter__(self) -> typing.Iterator[ParamName]: + def __iter__(self) -> Iterator[ParamName]: return iter(self.__data) def __len__(self) -> int: @@ -88,7 +88,7 @@ def __repr__(self) -> str: # pylint:disable=too-many-branches -def format_param_args(param: ParamName, details: Any) -> Union[ParamArg, List[ParamVal]]: +def format_param_args(param: ParamName, details: Any) -> Union[ParamArg, list[ParamVal]]: """Ensures each user-inputted parameter is a properly typed list. Also provides custom support for certain parameters.""" if not isinstance(details, list): @@ -127,7 +127,7 @@ def format_param_args(param: ParamName, details: Any) -> Union[ParamArg, List[Pa return details -def format_params(**params: Any) -> Dict[ParamName, Union[ParamArg, ParamVal]]: +def format_params(**params: Any) -> dict[ParamName, Union[ParamArg, ParamVal]]: """Converts params to a dictionary whose keys are parameter names and whose values are single ``ParamaterArg`` objects or lists of parameter values.""" return { diff --git a/pennylane/debugging/snapshot.py b/pennylane/debugging/snapshot.py index f58a2893477..df06ffa7dfd 100644 --- a/pennylane/debugging/snapshot.py +++ b/pennylane/debugging/snapshot.py @@ -16,12 +16,11 @@ """ import warnings from functools import partial -from typing import Callable, Sequence import pennylane as qml -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform -from pennylane.typing import Result, ResultBatch +from pennylane.typing import PostprocessingFn def _is_snapshot_compatible(dev): @@ -55,7 +54,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback): @transform -def snapshots(tape: QuantumTape) -> tuple[Sequence[QuantumTape], Callable[[ResultBatch], Result]]: +def snapshots(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""This transform processes :class:`~pennylane.Snapshot` instances contained in a circuit, depending on the compatibility of the execution device. For supported devices, the snapshots' measurements are computed as the execution progresses. diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py index f0750efbd5a..475a84aac0b 100644 --- a/pennylane/devices/__init__.py +++ b/pennylane/devices/__init__.py @@ -149,6 +149,7 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi """ from .execution_config import ExecutionConfig, DefaultExecutionConfig, MCMConfig +from .device_constructor import device, refresh_devices from .device_api import Device from .default_qubit import DefaultQubit @@ -163,5 +164,15 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi from .default_clifford import DefaultClifford from .default_tensor import DefaultTensor from .null_qubit import NullQubit +from .default_qutrit import DefaultQutrit from .default_qutrit_mixed import DefaultQutritMixed from .._device import Device as LegacyDevice +from .._device import DeviceError + + +# pylint: disable=undefined-variable +def __getattr__(name): + if name == "plugin_devices": + return device_constructor.plugin_devices + + raise AttributeError(f"module 'pennylane.devices' has no attribute '{name}'") diff --git a/pennylane/devices/default_clifford.py b/pennylane/devices/default_clifford.py index c5d23c50ff4..3585e1d7d8b 100644 --- a/pennylane/devices/default_clifford.py +++ b/pennylane/devices/default_clifford.py @@ -16,9 +16,10 @@ """ import concurrent.futures +from collections.abc import Sequence from dataclasses import replace from functools import partial -from typing import Sequence, Tuple, Union +from typing import Union import numpy as np @@ -39,7 +40,7 @@ VnEntropyMP, ) from pennylane.ops.qubit.observables import BasisStateProjector -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import convert_to_numpy_parameters from pennylane.transforms.core import TransformProgram from pennylane.typing import Result, ResultBatch @@ -64,8 +65,6 @@ except (ModuleNotFoundError, ImportError) as import_error: # pragma: no cover has_stim = False -Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTape_or_Batch = Union[QuantumTape, Sequence[QuantumTape]] # Updated observable list _OBSERVABLES_MAP = { @@ -451,7 +450,7 @@ def _setup_execution_config(self, execution_config: ExecutionConfig) -> Executio def preprocess( self, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Tuple[TransformProgram, ExecutionConfig]: + ) -> tuple[TransformProgram, ExecutionConfig]: """This function defines the device transform program to be applied and an updated device configuration. Args: @@ -499,9 +498,9 @@ def preprocess( def execute( self, - circuits: QuantumTape_or_Batch, + circuits: Union[QuantumTape, QuantumTapeBatch], execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Result_or_ResultBatch: + ) -> Union[Result, ResultBatch]: max_workers = execution_config.device_options.get("max_workers", self._max_workers) if max_workers is None: seeds = self._rng.integers(2**31 - 1, size=len(circuits)) diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py index aebbe4bbeeb..bcd97c003b8 100644 --- a/pennylane/devices/default_qubit.py +++ b/pennylane/devices/default_qubit.py @@ -20,7 +20,7 @@ from dataclasses import replace from functools import partial from numbers import Number -from typing import Callable, Optional, Sequence, Tuple, Union +from typing import Optional, Union import numpy as np @@ -28,10 +28,10 @@ from pennylane.logging import debug_logger, debug_logger_init from pennylane.measurements.mid_measure import MidMeasureMP from pennylane.ops.op_math.condition import Conditional -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import convert_to_numpy_parameters from pennylane.transforms.core import TransformProgram -from pennylane.typing import Result, ResultBatch +from pennylane.typing import PostprocessingFn, Result, ResultBatch from . import Device from .execution_config import DefaultExecutionConfig, ExecutionConfig @@ -53,11 +53,7 @@ logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) -Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTapeBatch = Sequence[QuantumTape] QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch] -# always a function from a resultbatch to either a result or a result batch -PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch] def stopping_condition(op: qml.operation.Operator) -> bool: @@ -163,7 +159,7 @@ def all_state_postprocessing(results, measurements, wire_order): @qml.transform def adjoint_state_measurements( tape: QuantumTape, device_vjp=False -) -> Tuple[Tuple[QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Perform adjoint measurement preprocessing. * Allows a tape with only expectation values through unmodified @@ -516,7 +512,7 @@ def supports_derivatives( def preprocess( self, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Tuple[TransformProgram, ExecutionConfig]: + ) -> tuple[TransformProgram, ExecutionConfig]: """This function defines the device transform program to be applied and an updated device configuration. Args: @@ -615,7 +611,7 @@ def execute( self, circuits: QuantumTape_or_Batch, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Result_or_ResultBatch: + ) -> Union[Result, ResultBatch]: self.reset_prng_key() max_workers = execution_config.device_options.get("max_workers", self._max_workers) @@ -732,7 +728,7 @@ def supports_jvp( def compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): max_workers = execution_config.device_options.get("max_workers", self._max_workers) @@ -752,7 +748,7 @@ def compute_jvp( def execute_and_compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): self.reset_prng_key() @@ -800,7 +796,7 @@ def supports_vjp( def compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): r"""The vector jacobian product used in reverse-mode differentiation. ``DefaultQubit`` uses the @@ -868,7 +864,7 @@ def _state(circuit): def execute_and_compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): self.reset_prng_key() diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py index 5ad6b629625..abcc6e0452f 100644 --- a/pennylane/devices/default_qubit_autograd.py +++ b/pennylane/devices/default_qubit_autograd.py @@ -14,12 +14,15 @@ """This module contains an autograd implementation of the :class:`~.DefaultQubitLegacy` reference plugin. """ +import warnings + +from pennylane import PennyLaneDeprecationWarning from pennylane import numpy as pnp from pennylane.devices import DefaultQubitLegacy class DefaultQubitAutograd(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd. **Short name:** ``default.qubit.autograd`` @@ -34,6 +37,9 @@ class DefaultQubitAutograd(DefaultQubitLegacy): pip install autograd + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.autograd`` is designed to be used with end-to-end classical backpropagation @@ -104,6 +110,14 @@ def _const_mul(constant, array): return constant * array def __init__(self, wires, *, shots=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + PennyLaneDeprecationWarning, + ) + r_dtype = pnp.float64 c_dtype = pnp.complex128 super().__init__(wires, shots=shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic) diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py index f349b8ff7c9..0cae4827e41 100644 --- a/pennylane/devices/default_qubit_jax.py +++ b/pennylane/devices/default_qubit_jax.py @@ -14,6 +14,8 @@ """This module contains a jax implementation of the :class:`~.DefaultQubitLegacy` reference plugin. """ +import warnings + # pylint: disable=ungrouped-imports import numpy as np @@ -34,7 +36,7 @@ class DefaultQubitJax(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using jax. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using jax. **Short name:** ``default.qubit.jax`` @@ -49,6 +51,9 @@ class DefaultQubitJax(DefaultQubitLegacy): pip install jax jaxlib + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.jax`` device is designed to be used with end-to-end classical backpropagation @@ -165,6 +170,14 @@ def circuit(): operations = DefaultQubitLegacy.operations.union({"ParametrizedEvolution"}) def __init__(self, wires, *, shots=None, prng_key=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + if jax.config.read("jax_enable_x64"): c_dtype = jnp.complex128 r_dtype = jnp.float64 diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py index 426a0fe31d4..f351b2aa455 100644 --- a/pennylane/devices/default_qubit_legacy.py +++ b/pennylane/devices/default_qubit_legacy.py @@ -20,8 +20,8 @@ """ import functools import itertools +import warnings from string import ascii_letters as ABC -from typing import List import numpy as np from scipy.sparse import csr_matrix @@ -78,13 +78,13 @@ def _get_slice(index, axis, num_axes): # pylint: disable=unused-argument class DefaultQubitLegacy(QubitDevice): - """Default qubit device for PennyLane. + r"""Default qubit device for PennyLane. .. warning:: - This is the legacy implementation of DefaultQubit. It has been replaced by - ``qml.devices.DefaultQubit``, which can be accessed with the familiar constructor, - ``qml.device("default.qubit")``. + This is the legacy implementation of DefaultQubit and is deprecated. It has been replaced by + :class:`~pennylane.devices.DefaultQubit`, which can be accessed with the familiar constructor, + ``qml.device("default.qubit")``, and now supports backpropagation. This change will not alter device behaviour for most workflows, but may have implications for plugin developers and users who directly interact with device methods. Please consult @@ -207,6 +207,14 @@ class DefaultQubitLegacy(QubitDevice): def __init__( self, wires, *, r_dtype=np.float64, c_dtype=np.complex128, shots=None, analytic=None ): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + super().__init__(wires, shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic) self._debugger = None @@ -1090,7 +1098,7 @@ def classical_shadow(self, obs, circuit): return self._cast(self._stack([outcomes, recipes]), dtype=np.int8) - def _get_diagonalizing_gates(self, circuit: qml.tape.QuantumTape) -> List[Operation]: + def _get_diagonalizing_gates(self, circuit: qml.tape.QuantumTape) -> list[Operation]: meas_filtered = [ m for m in circuit.measurements diff --git a/pennylane/devices/default_qubit_tf.py b/pennylane/devices/default_qubit_tf.py index 48d7c60b788..2171aa7592e 100644 --- a/pennylane/devices/default_qubit_tf.py +++ b/pennylane/devices/default_qubit_tf.py @@ -15,6 +15,7 @@ reference plugin. """ import itertools +import warnings import numpy as np from packaging.version import Version @@ -42,7 +43,7 @@ class DefaultQubitTF(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using TensorFlow. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using TensorFlow. **Short name:** ``default.qubit.tf`` @@ -57,6 +58,9 @@ class DefaultQubitTF(DefaultQubitLegacy): pip install tensorflow>=2.0 + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.tf`` is designed to be used with end-to-end classical backpropagation @@ -162,6 +166,14 @@ def _asarray(array, dtype=None): return res def __init__(self, wires, *, shots=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + r_dtype = tf.float64 c_dtype = tf.complex128 diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py index e0bb276e513..3ed26025b60 100644 --- a/pennylane/devices/default_qubit_torch.py +++ b/pennylane/devices/default_qubit_torch.py @@ -34,6 +34,7 @@ import numpy as np +from pennylane import PennyLaneDeprecationWarning from pennylane.ops.qubit.attributes import diagonal_in_z_basis from . import DefaultQubitLegacy @@ -43,7 +44,7 @@ class DefaultQubitTorch(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using PyTorch. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using PyTorch. **Short name:** ``default.qubit.torch`` @@ -58,6 +59,10 @@ class DefaultQubitTorch(DefaultQubitLegacy): pip install torch>=1.8.0 + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + + **Example** The ``default.qubit.torch`` is designed to be used with end-to-end classical backpropagation @@ -165,6 +170,13 @@ def circuit(x): _ndim = staticmethod(lambda tensor: tensor.ndim) def __init__(self, wires, *, shots=None, analytic=None, torch_device=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + PennyLaneDeprecationWarning, + ) # Store if the user specified a Torch device. Otherwise the execute # method attempts to infer the Torch device from the gate parameters. self._torch_device_specified = torch_device is not None diff --git a/pennylane/devices/default_qutrit_mixed.py b/pennylane/devices/default_qutrit_mixed.py index cdad2464754..bf8ae1eea7d 100644 --- a/pennylane/devices/default_qutrit_mixed.py +++ b/pennylane/devices/default_qutrit_mixed.py @@ -15,16 +15,17 @@ computations.""" import logging import warnings +from collections.abc import Callable, Sequence from dataclasses import replace from functools import partial -from typing import Callable, Optional, Sequence, Tuple, Union +from typing import Optional, Union import numpy as np import pennylane as qml from pennylane.logging import debug_logger, debug_logger_init from pennylane.ops import _qutrit__channel__ops__ as channels -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.core import TransformProgram from pennylane.typing import Result, ResultBatch @@ -46,11 +47,8 @@ logger.addHandler(logging.NullHandler()) Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTapeBatch = Sequence[QuantumTape] QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch] -# always a function from a resultbatch to either a result or a result batch -PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch] observables = { "THermitian", @@ -342,7 +340,7 @@ def _setup_execution_config(self, execution_config: ExecutionConfig) -> Executio def preprocess( self, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Tuple[TransformProgram, ExecutionConfig]: + ) -> tuple[TransformProgram, ExecutionConfig]: """This function defines the device transform program to be applied and an updated device configuration. diff --git a/pennylane/devices/default_tensor.py b/pennylane/devices/default_tensor.py index 9381ac62b49..78c9046c466 100644 --- a/pennylane/devices/default_tensor.py +++ b/pennylane/devices/default_tensor.py @@ -17,10 +17,11 @@ # pylint: disable=protected-access import copy import warnings +from collections.abc import Callable from dataclasses import replace from functools import singledispatch from numbers import Number -from typing import Callable, Optional, Sequence, Tuple, Union +from typing import Optional, Union import numpy as np @@ -42,16 +43,13 @@ ) from pennylane.operation import Observable, Operation, Tensor from pennylane.ops import LinearCombination, Prod, SProd, Sum -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch from pennylane.templates.subroutines.trotter import _recursive_expression from pennylane.transforms.core import TransformProgram from pennylane.typing import Result, ResultBatch, TensorLike from pennylane.wires import WireError -Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTapeBatch = Sequence[QuantumTape] QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch] -PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch] has_quimb = True @@ -628,7 +626,7 @@ def execute( self, circuits: QuantumTape_or_Batch, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Result_or_ResultBatch: + ) -> Union[Result, ResultBatch]: """Execute a circuit or a batch of circuits and turn it into results. Args: @@ -885,7 +883,7 @@ def supports_vjp( def compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): r"""The vector-Jacobian product used in reverse-mode differentiation. @@ -907,7 +905,7 @@ def compute_vjp( def execute_and_compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): """Calculate both the results and the vector-Jacobian product used in reverse-mode differentiation. diff --git a/pennylane/devices/device_api.py b/pennylane/devices/device_api.py index b472973d369..4234a6dbd4c 100644 --- a/pennylane/devices/device_api.py +++ b/pennylane/devices/device_api.py @@ -19,21 +19,18 @@ from collections.abc import Iterable from dataclasses import replace from numbers import Number -from typing import Callable, Optional, Sequence, Tuple, Union +from typing import Optional, Union from pennylane import Tracker from pennylane.measurements import Shots -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.core import TransformProgram from pennylane.typing import Result, ResultBatch from pennylane.wires import Wires from .execution_config import DefaultExecutionConfig, ExecutionConfig -Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTapeBatch = Sequence[QuantumTape] QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch] -PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch] # pylint: disable=unused-argument, no-self-use @@ -220,7 +217,7 @@ def wires(self) -> Wires: def preprocess( self, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Tuple[TransformProgram, ExecutionConfig]: + ) -> tuple[TransformProgram, ExecutionConfig]: """Device preprocessing function. .. warning:: @@ -258,8 +255,11 @@ def preprocess( .. code-block:: python + from pennylane.tape import TapeBatch + from pennylane.typing import PostprocessingFn + @transform - def my_preprocessing_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], callable): + def my_preprocessing_transform(tape: qml.tape.QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: # e.g. valid the measurements, expand the tape for the hardware execution, ... def blank_processing_fn(results): @@ -322,7 +322,7 @@ def execute( self, circuits: QuantumTape_or_Batch, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Result_or_ResultBatch: + ) -> Union[Result, ResultBatch]: """Execute a circuit or a batch of circuits and turn it into results. Args: @@ -543,7 +543,7 @@ def execute_and_compute_derivatives( def compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): r"""The jacobian vector product used in forward mode calculation of derivatives. @@ -582,7 +582,7 @@ def compute_jvp( def execute_and_compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): """Execute a batch of circuits and compute their jacobian vector products. @@ -620,7 +620,7 @@ def supports_jvp( def compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): r"""The vector jacobian product used in reverse-mode differentiation. @@ -660,7 +660,7 @@ def compute_vjp( def execute_and_compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number, ...], execution_config: ExecutionConfig = DefaultExecutionConfig, ): r"""Calculate both the results and the vector jacobian product used in reverse-mode differentiation. diff --git a/pennylane/devices/device_constructor.py b/pennylane/devices/device_constructor.py new file mode 100644 index 00000000000..107243f7718 --- /dev/null +++ b/pennylane/devices/device_constructor.py @@ -0,0 +1,291 @@ +# Copyright 2018-2024 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This module contains code for the main device construction delegation logic. +""" +import warnings +from importlib import metadata +from sys import version_info + +from semantic_version import SimpleSpec, Version + +import pennylane as qml + + +def _get_device_entrypoints(): + """Returns a dictionary mapping the device short name to the + loadable entrypoint""" + + entries = ( + metadata.entry_points()["pennylane.plugins"] + if version_info[:2] == (3, 9) + # pylint:disable=unexpected-keyword-arg + else metadata.entry_points(group="pennylane.plugins") + ) + return {entry.name: entry for entry in entries} + + +# get list of installed devices +plugin_devices = _get_device_entrypoints() + + +def refresh_devices(): + """Scan installed PennyLane plugins to refresh the device list.""" + + # This function does not return anything; instead, it has a side effect + # which is to update the global plugin_devices variable. + + # We wish to retain the behaviour of a global plugin_devices dictionary, + # as re-importing metadata can be a very slow operation on systems + # with a large number of installed packages. + + global plugin_devices # pylint:disable=global-statement + plugin_devices = _get_device_entrypoints() + + +# pylint: disable=protected-access +def device(name, *args, **kwargs): + r""" + Load a device and return the instance. + + This function is used to load a particular quantum device, + which can then be used to construct QNodes. + + PennyLane comes with support for the following devices: + + * :mod:`'default.qubit' `: a simple + state simulator of qubit-based quantum circuit architectures. + + * :mod:`'default.mixed' `: a mixed-state + simulator of qubit-based quantum circuit architectures. + + * ``'lightning.qubit'``: a more performant state simulator of qubit-based + quantum circuit architectures written in C++. + + * :mod:`'default.qutrit' `: a simple + state simulator of qutrit-based quantum circuit architectures. + + * :mod:`'default.qutrit.mixed' `: a + mixed-state simulator of qutrit-based quantum circuit architectures. + + * :mod:`'default.gaussian' `: a simple simulator + of Gaussian states and operations on continuous-variable circuit architectures. + + * :mod:`'default.clifford' `: an efficient + simulator of Clifford circuits. + + * :mod:`'default.tensor' `: a simulator + of quantum circuits based on tensor networks. + + Additional devices are supported through plugins β€” see + the `available plugins `_ for more + details. To list all currently installed devices, run + :func:`qml.about `. + + Args: + name (str): the name of the device to load + wires (int): the number of wires (subsystems) to initialise + the device with. Note that this is optional for certain + devices, such as ``default.qubit`` + + Keyword Args: + config (pennylane.Configuration): a PennyLane configuration object + that contains global and/or device specific configurations. + custom_decomps (Dict[Union(str, Operator), Callable]): Custom + decompositions to be applied by the device at runtime. + decomp_depth (int): For when custom decompositions are specified, + the maximum expansion depth used by the expansion function. + + .. warning:: + + The ``decomp_depth`` argument is deprecated and will be removed in version 0.39. + + All devices must be loaded by specifying their **short-name** as listed above, + followed by the **wires** (subsystems) you wish to initialize. The ``wires`` + argument can be an integer, in which case the wires of the device are addressed + by consecutive integers: + + .. code-block:: python + + dev = qml.device('default.qubit', wires=5) + + def circuit(): + qml.Hadamard(wires=1) + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[3, 4]) + ... + + The ``wires`` argument can also be a sequence of unique numbers or strings, specifying custom wire labels + that the user employs to address the wires: + + .. code-block:: python + + dev = qml.device('default.qubit', wires=['ancilla', 'q11', 'q12', -1, 1]) + + def circuit(): + qml.Hadamard(wires='q11') + qml.Hadamard(wires=['ancilla']) + qml.CNOT(wires=['q12', -1]) + ... + + On some newer devices, such as ``default.qubit``, the ``wires`` argument can be omitted altogether, + and instead the wires will be computed when executing a circuit depending on its contents. + + >>> dev = qml.device("default.qubit") + + Most devices accept a ``shots`` argument which specifies how many circuit executions + are used to estimate stochastic return values. As an example, ``qml.sample()`` measurements + will return as many samples as specified in the shots argument. The shots argument can be + changed on a per-call basis using the built-in ``shots`` keyword argument. Note that the + ``shots`` argument can be a single integer or a list of shot values. + + .. code-block:: python + + dev = qml.device('default.qubit', wires=1, shots=10) + + @qml.qnode(dev) + def circuit(a): + qml.RX(a, wires=0) + return qml.sample(qml.Z(0)) + + >>> circuit(0.8) # 10 samples are returned + array([ 1, 1, 1, 1, -1, 1, 1, -1, 1, 1]) + >>> circuit(0.8, shots=[3, 4, 4]) # default is overwritten for this call + (array([1, 1, 1]), array([ 1, -1, 1, 1]), array([1, 1, 1, 1])) + >>> circuit(0.8) # back to default of 10 samples + array([ 1, -1, 1, 1, -1, 1, 1, 1, 1, 1]) + + When constructing a device, we may optionally pass a dictionary of custom + decompositions to be applied to certain operations upon device execution. + This is useful for enabling support of gates on devices where they would normally + be unsupported. + + For example, suppose we are running on an ion trap device which does not + natively implement the CNOT gate, but we would still like to write our + circuits in terms of CNOTs. On a ion trap device, CNOT can be implemented + using the ``IsingXX`` gate. We first define a decomposition function + (such functions have the signature ``decomposition(*params, wires)``): + + .. code-block:: python + + def ion_trap_cnot(wires, **_): + return [ + qml.RY(np.pi/2, wires=wires[0]), + qml.IsingXX(np.pi/2, wires=wires), + qml.RX(-np.pi/2, wires=wires[0]), + qml.RY(-np.pi/2, wires=wires[0]), + qml.RY(-np.pi/2, wires=wires[1]) + ] + + Next, we create a device, and a QNode for testing. When constructing the + QNode, we can set the expansion strategy to ``"device"`` to ensure the + decomposition is applied and will be viewable when we draw the circuit. + Note that custom decompositions should accept keyword arguments even when + it is not used. + + .. code-block:: python + + # As the CNOT gate normally has no decomposition, we can use default.qubit + # here for expository purposes. + dev = qml.device( + 'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot} + ) + + @qml.qnode(dev, expansion_strategy="device") + def run_cnot(): + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.X(1)) + + >>> print(qml.draw(run_cnot)()) + 0: ──RY(1.57)─╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)── + 1: ───────────╰IsingXX(1.57)──RY(-1.57)───────────── + + Some devices may accept additional arguments. For instance, + ``default.gaussian`` accepts the keyword argument ``hbar``, to set + the convention used in the commutation relation :math:`[\x,\p]=i\hbar` + (by default set to 2). + + Please refer to the documentation for the individual devices to see any + additional arguments that might be required or supported. + """ + if name not in plugin_devices: + # Device does not exist in the loaded device list. + # Attempt to refresh the devices, in case the user + # installed the plugin during the current Python session. + refresh_devices() + + if name in plugin_devices: + options = {} + + # load global configuration settings if available + config = kwargs.get("config", qml.default_config) + + if config: + # combine configuration options with keyword arguments. + # Keyword arguments take preference, followed by device options, + # followed by plugin options, followed by global options. + options.update(config["main"]) + options.update(config[name.split(".")[0] + ".global"]) + options.update(config[name]) + + # Pop the custom decomposition keyword argument; we will use it here + # only and not pass it to the device. + custom_decomps = kwargs.pop("custom_decomps", None) + decomp_depth = kwargs.pop("decomp_depth", None) + + if decomp_depth is not None: + warnings.warn( + "The decomp_depth argument is deprecated and will be removed in version 0.39. ", + qml.PennyLaneDeprecationWarning, + ) + else: + decomp_depth = 10 + + kwargs.pop("config", None) + options.update(kwargs) + + # loads the device class + plugin_device_class = plugin_devices[name].load() + + if hasattr(plugin_device_class, "pennylane_requires") and Version( + qml.version() + ) not in SimpleSpec(plugin_device_class.pennylane_requires): + raise qml.DeviceError( + f"The {name} plugin requires PennyLane versions {plugin_device_class.pennylane_requires}, " + f"however PennyLane version {qml.version()} is installed." + ) + + # Construct the device + dev = plugin_device_class(*args, **options) + + # Once the device is constructed, we set its custom expansion function if + # any custom decompositions were specified. + + if custom_decomps is not None: + if isinstance(dev, qml.devices.LegacyDevice): + custom_decomp_expand_fn = qml.transforms.create_decomp_expand_fn( + custom_decomps, dev, decomp_depth=decomp_depth + ) + dev.custom_expand(custom_decomp_expand_fn) + else: + custom_decomp_preprocess = qml.transforms.tape_expand._create_decomp_preprocessing( + custom_decomps, dev, decomp_depth=decomp_depth + ) + dev.preprocess = custom_decomp_preprocess + + return dev + + raise qml.DeviceError( + f"Device {name} does not exist. Make sure the required plugin is installed." + ) diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py index 9204fe1c51e..9105edc2cee 100644 --- a/pennylane/devices/null_qubit.py +++ b/pennylane/devices/null_qubit.py @@ -22,7 +22,7 @@ from dataclasses import replace from functools import singledispatch from numbers import Number -from typing import Callable, Sequence, Tuple, Union +from typing import Union import numpy as np @@ -40,7 +40,7 @@ Shots, StateMP, ) -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.core import TransformProgram from pennylane.typing import Result, ResultBatch @@ -51,11 +51,7 @@ logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) -Result_or_ResultBatch = Union[Result, ResultBatch] -QuantumTapeBatch = Sequence[QuantumTape] QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch] -# always a function from a resultbatch to either a result or a result batch -PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch] @singledispatch @@ -282,7 +278,7 @@ def _max_workers(self): # pylint: disable=cell-var-from-loop def preprocess( self, execution_config=DefaultExecutionConfig - ) -> Tuple[TransformProgram, ExecutionConfig]: + ) -> tuple[TransformProgram, ExecutionConfig]: program, _ = DefaultQubit.preprocess(self, execution_config) for t in program: if t.transform == decompose.transform: @@ -323,7 +319,7 @@ def execute( self, circuits: QuantumTape_or_Batch, execution_config: ExecutionConfig = DefaultExecutionConfig, - ) -> Result_or_ResultBatch: + ) -> Union[Result, ResultBatch]: if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug( """Entry with args=(circuits=%s) called by=%s""", @@ -384,7 +380,7 @@ def execute_and_compute_derivatives( def compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number], execution_config: ExecutionConfig = DefaultExecutionConfig, ): return tuple(self._jvp(c, INTERFACE_TO_LIKE[execution_config.interface]) for c in circuits) @@ -392,7 +388,7 @@ def compute_jvp( def execute_and_compute_jvp( self, circuits: QuantumTape_or_Batch, - tangents: Tuple[Number], + tangents: tuple[Number], execution_config: ExecutionConfig = DefaultExecutionConfig, ): results = tuple( @@ -405,7 +401,7 @@ def execute_and_compute_jvp( def compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number], execution_config: ExecutionConfig = DefaultExecutionConfig, ): return tuple(self._vjp(c, INTERFACE_TO_LIKE[execution_config.interface]) for c in circuits) @@ -413,7 +409,7 @@ def compute_vjp( def execute_and_compute_vjp( self, circuits: QuantumTape_or_Batch, - cotangents: Tuple[Number], + cotangents: tuple[Number], execution_config: ExecutionConfig = DefaultExecutionConfig, ): results = tuple( diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py index 2f38f1f9a54..ec7d337bc9e 100644 --- a/pennylane/devices/preprocess.py +++ b/pennylane/devices/preprocess.py @@ -18,21 +18,21 @@ import os import warnings +from collections.abc import Callable, Generator, Sequence from copy import copy from itertools import chain -from typing import Callable, Generator, Optional, Sequence, Union +from typing import Optional, Union import pennylane as qml from pennylane import DeviceError, Snapshot, transform from pennylane.measurements import SampleMeasurement, StateMeasurement from pennylane.operation import StatePrepBase, Tensor -from pennylane.typing import Result, ResultBatch +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from pennylane.wires import WireError from .execution_config import MCMConfig -PostprocessingFn = Callable[[ResultBatch], Union[Result, ResultBatch]] - def null_postprocessing(results): """A postprocessing function returned by a transform that only converts the batch of results @@ -83,7 +83,7 @@ def _operator_decomposition_gen( @transform def no_sampling( tape: qml.tape.QuantumTape, name: str = "device" -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Raises an error if the tape has finite shots. Args: @@ -107,7 +107,7 @@ def no_sampling( @transform def validate_device_wires( tape: qml.tape.QuantumTape, wires: Optional[qml.wires.Wires] = None, name: str = "device" -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Validates that all wires present in the tape are in the set of provided wires. Adds the device wires to measurement processes like :class:`~.measurements.StateMP` that are broadcasted across all available wires. @@ -152,7 +152,7 @@ def mid_circuit_measurements( device, mcm_config=MCMConfig(), interface=None, -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Provide the transform to handle mid-circuit measurements. If the tape or device uses finite-shot, use the native implementation (i.e. no transform), @@ -175,7 +175,7 @@ def mid_circuit_measurements( @transform def validate_multiprocessing_workers( tape: qml.tape.QuantumTape, max_workers: int, device -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Validates the number of workers for multiprocessing. Checks that the CPU is not oversubscribed and warns user if it is, @@ -234,7 +234,7 @@ def validate_multiprocessing_workers( @transform def validate_adjoint_trainable_params( tape: qml.tape.QuantumTape, -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Raises a warning if any of the observables is trainable, and raises an error if any trainable parameters belong to state-prep operations. Can be used in validating circuits for adjoint differentiation. @@ -270,7 +270,7 @@ def decompose( max_expansion: Union[int, None] = None, name: str = "device", error: Exception = DeviceError, -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Decompose operations until the stopping condition is met. Args: @@ -386,7 +386,7 @@ def validate_observables( tape: qml.tape.QuantumTape, stopping_condition: Callable[[qml.operation.Operator], bool], name: str = "device", -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Validates the observables and measurements for a circuit. Args: @@ -428,7 +428,7 @@ def validate_observables( @transform def validate_measurements( tape: qml.tape.QuantumTape, analytic_measurements=None, sample_measurements=None, name="device" -) -> tuple[Sequence[qml.tape.QuantumTape], Callable[[ResultBatch], Result]]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Validates the supported state and sample based measurement processes. Args: diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py index 7f25cb3a858..855a2d81f5d 100644 --- a/pennylane/devices/qubit/adjoint_jacobian.py +++ b/pennylane/devices/qubit/adjoint_jacobian.py @@ -14,7 +14,6 @@ """Functions to apply adjoint jacobian differentiation""" import logging from numbers import Number -from typing import Tuple import numpy as np @@ -150,7 +149,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None): @debug_logger -def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None): +def adjoint_jvp(tape: QuantumTape, tangents: tuple[Number], state=None): """The jacobian vector product used in forward mode calculation of derivatives. Implements the adjoint method outlined in @@ -324,7 +323,7 @@ def _get_vjp_bras(tape, cotangents, ket): @debug_logger -def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None): +def adjoint_vjp(tape: QuantumTape, cotangents: tuple[Number, ...], state=None): """The vector jacobian product used in reverse-mode differentiation. Implements the adjoint method outlined in diff --git a/pennylane/devices/qubit/initialize_state.py b/pennylane/devices/qubit/initialize_state.py index e2cf9e6e780..86101deb8e8 100644 --- a/pennylane/devices/qubit/initialize_state.py +++ b/pennylane/devices/qubit/initialize_state.py @@ -13,7 +13,8 @@ # limitations under the License. """Functions to prepare a state.""" -from typing import Iterable, Union +from collections.abc import Iterable +from typing import Union import numpy as np diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py index 688ca50b001..ae2ddb07c02 100644 --- a/pennylane/devices/qubit/measure.py +++ b/pennylane/devices/qubit/measure.py @@ -14,7 +14,7 @@ """ Code relevant for performing measurements on a state. """ -from typing import Callable +from collections.abc import Callable from scipy.sparse import csr_matrix diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py index be2a2ae8b14..2d87d3e504e 100644 --- a/pennylane/devices/qubit/sampling.py +++ b/pennylane/devices/qubit/sampling.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Functions to sample a state.""" -from typing import List, Tuple, Union +from typing import Union import numpy as np @@ -42,7 +42,7 @@ def jax_random_split(prng_key, num: int = 2): return split(prng_key, num=num) -def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]): +def _group_measurements(mps: list[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]): """ Group the measurements such that: - measurements with pauli observables pairwise-commute in each group @@ -146,7 +146,7 @@ def _get_num_executions_for_sum(obs): # pylint: disable=no-member -def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]: +def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> tuple[int, int]: """Get the total number of qpu executions and shots. Args: @@ -190,7 +190,7 @@ def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]: def _apply_diagonalizing_gates( - mps: List[SampleMeasurement], state: np.ndarray, is_state_batched: bool = False + mps: list[SampleMeasurement], state: np.ndarray, is_state_batched: bool = False ): if len(mps) == 1: diagonalizing_gates = mps[0].diagonalizing_gates() @@ -207,14 +207,14 @@ def _apply_diagonalizing_gates( # pylint:disable = too-many-arguments def measure_with_samples( - measurements: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]], + measurements: list[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]], state: np.ndarray, shots: Shots, is_state_batched: bool = False, rng=None, prng_key=None, mid_measurements: dict = None, -) -> List[TensorLike]: +) -> list[TensorLike]: """ Returns the samples of the measurement process performed on the given state. This function assumes that the user-defined wire labels in the measurement process @@ -281,7 +281,7 @@ def measure_with_samples( def _measure_with_samples_diagonalizing_gates( - mps: List[SampleMeasurement], + mps: list[SampleMeasurement], state: np.ndarray, shots: Shots, is_state_batched: bool = False, @@ -351,7 +351,7 @@ def _process_single_shot(samples): def _measure_classical_shadow( - mp: List[Union[ClassicalShadowMP, ShadowExpvalMP]], + mp: list[Union[ClassicalShadowMP, ShadowExpvalMP]], state: np.ndarray, shots: Shots, is_state_batched: bool = False, @@ -390,7 +390,7 @@ def _measure_classical_shadow( def _measure_hamiltonian_with_samples( - mp: List[SampleMeasurement], + mp: list[SampleMeasurement], state: np.ndarray, shots: Shots, is_state_batched: bool = False, @@ -421,7 +421,7 @@ def _sum_for_single_shot(s, prng_key=None): def _measure_sum_with_samples( - mp: List[SampleMeasurement], + mp: list[SampleMeasurement], state: np.ndarray, shots: Shots, is_state_batched: bool = False, diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py index e55741b029d..9017e300316 100644 --- a/pennylane/devices/qubit/simulate.py +++ b/pennylane/devices/qubit/simulate.py @@ -230,7 +230,6 @@ def measure_final_state(circuit, state, is_state_batched, **execution_kwargs) -> mid_measurements = execution_kwargs.get("mid_measurements", None) # analytic case - if not circuit.shots: if mid_measurements is not None: raise TypeError("Native mid-circuit measurements are only supported with finite shots.") @@ -243,7 +242,6 @@ def measure_final_state(circuit, state, is_state_batched, **execution_kwargs) -> ) # finite-shot case - rng = default_rng(rng) results = measure_with_samples( circuit.measurements, diff --git a/pennylane/devices/qutrit_mixed/initialize_state.py b/pennylane/devices/qutrit_mixed/initialize_state.py index 7906e0894dc..0b097097e22 100644 --- a/pennylane/devices/qutrit_mixed/initialize_state.py +++ b/pennylane/devices/qutrit_mixed/initialize_state.py @@ -13,7 +13,8 @@ # limitations under the License. """Functions to prepare a qutrit mixed state.""" -from typing import Iterable, Union +from collections.abc import Iterable +from typing import Union import pennylane as qml from pennylane.operation import StatePrepBase diff --git a/pennylane/devices/qutrit_mixed/measure.py b/pennylane/devices/qutrit_mixed/measure.py index d93a0eeff79..3489cb3ee35 100644 --- a/pennylane/devices/qutrit_mixed/measure.py +++ b/pennylane/devices/qutrit_mixed/measure.py @@ -15,8 +15,8 @@ Code relevant for performing measurements on a qutrit mixed state. """ +from collections.abc import Callable from string import ascii_letters as alphabet -from typing import Callable from pennylane import math, queuing from pennylane.measurements import ( diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py index 50e692bc6b7..83874185ca2 100644 --- a/pennylane/drawer/draw.py +++ b/pennylane/drawer/draw.py @@ -20,25 +20,30 @@ from functools import wraps import pennylane as qml +from pennylane.data.base.typing_util import UNSET from .tape_mpl import tape_mpl from .tape_text import tape_text -_level_sentinel = object() - def _determine_draw_level(kwargs, qnode=None): - sentinel = _level_sentinel - - level = kwargs.get("level", sentinel) - expansion_strategy = kwargs.get("expansion_strategy", sentinel) + level = kwargs.get("level", UNSET) + expansion_strategy = kwargs.get("expansion_strategy", UNSET) - if all(val != sentinel for val in (level, expansion_strategy)): + if all(val != UNSET for val in (level, expansion_strategy)): raise ValueError("Either 'level' or 'expansion_strategy' need to be set, but not both.") - if level == sentinel: - if expansion_strategy == sentinel: - return qnode.expansion_strategy if qnode else sentinel + if expansion_strategy != UNSET: + warnings.warn( + "The 'expansion_strategy' argument is deprecated and will be removed in " + "version 0.39. Instead, use the 'level' argument which offers more flexibility " + "and options.", + qml.PennyLaneDeprecationWarning, + ) + + if level == UNSET: + if expansion_strategy == UNSET: + return qnode.expansion_strategy if qnode else UNSET return expansion_strategy return level @@ -95,6 +100,9 @@ def draw( as it allows for the same values as ``expansion_strategy`` and offers more flexibility in choosing the desired transforms/expansions. + .. warning:: + The ``expansion_strategy`` argument is deprecated and will be removed in version 0.39. Use the ``level`` + argument instead to specify the resulting tape you want. **Example** @@ -173,7 +181,7 @@ def longer_circuit(params): qml.StronglyEntanglingLayers(params, wires=range(3)) return [qml.expval(qml.Z(i)) for i in range(3)] - >>> print(qml.draw(longer_circuit, max_length=60, expansion_strategy="device")(params)) + >>> print(qml.draw(longer_circuit, max_length=60, level="device")(params)) 0: ──Rot(0.77,0.44,0.86)─╭●────╭X──Rot(0.45,0.37,0.93)─╭●─╭X 1: ──Rot(0.70,0.09,0.98)─╰X─╭●─│───Rot(0.64,0.82,0.44)─│──╰● 2: ──Rot(0.76,0.79,0.13)────╰X─╰●──Rot(0.23,0.55,0.06)─╰X─── @@ -293,7 +301,7 @@ def circ(weights, order): level=_determine_draw_level(kwargs, qnode), ) - if _determine_draw_level(kwargs) != _level_sentinel: + if _determine_draw_level(kwargs) != UNSET: warnings.warn( "When the input to qml.draw is not a QNode, the expansion_strategy and level arguments are ignored.", UserWarning, @@ -443,6 +451,10 @@ def draw_mpl( as it allows for the same values as ``expansion_strategy`` and offers more flexibility in choosing the desired transforms/expansions. + .. warning:: + The ``expansion_strategy`` argument is deprecated and will be removed in version 0.39. Use the ``level`` + argument instead to specify the resulting tape you want. + .. warning:: Unlike :func:`~.draw`, this function can not draw the full result of a tape-splitting transform. In such cases, @@ -701,7 +713,7 @@ def circ(): **kwargs, ) - if _determine_draw_level(kwargs) != _level_sentinel: + if _determine_draw_level(kwargs) != UNSET: warnings.warn( "When the input to qml.draw is not a QNode, the expansion_strategy and level arguments are ignored.", UserWarning, diff --git a/pennylane/drawer/mpldrawer.py b/pennylane/drawer/mpldrawer.py index 75684f5bd67..ea2f2425d06 100644 --- a/pennylane/drawer/mpldrawer.py +++ b/pennylane/drawer/mpldrawer.py @@ -15,8 +15,7 @@ This module contains the MPLDrawer class for creating circuit diagrams with matplotlib """ import warnings -from collections.abc import Iterable -from typing import Sequence +from collections.abc import Iterable, Sequence has_mpl = True try: diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py index d73e4bcc68f..2626cde639d 100644 --- a/pennylane/drawer/tape_mpl.py +++ b/pennylane/drawer/tape_mpl.py @@ -459,6 +459,7 @@ def tape_mpl( if update_style := (has_mpl and style != "rcParams"): restore_params = mpl.rcParams.copy() _set_style(style) + try: return _tape_mpl( tape, diff --git a/pennylane/fermi/fermionic.py b/pennylane/fermi/fermionic.py index 01b5049e40f..73d8935cebe 100644 --- a/pennylane/fermi/fermionic.py +++ b/pennylane/fermi/fermionic.py @@ -273,6 +273,36 @@ def __pow__(self, value): return operator + def to_mat(self, n_orbitals=None): + r"""Return the matrix representation. + + Args: + n_orbitals (int or None): Number of orbitals. If not provided, it will be inferred from + the largest orbital index in the Fermi operator. + + Returns: + NumpyArray: Matrix representation of the :class:`~.FermiWord`. + + **Example** + + >>> w = FermiWord({(0, 0): '+', (1, 1): '-'}) + >>> w.to_mat() + array([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], + [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j], + [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]) + """ + largest_orb_id = max(key[1] for key in self.keys()) + 1 + if n_orbitals and n_orbitals < largest_orb_id: + raise ValueError( + f"n_orbitals cannot be smaller than {largest_orb_id}, got: {n_orbitals}." + ) + + largest_order = n_orbitals or largest_orb_id + mat = qml.jordan_wigner(self, ps=True).to_mat(wire_order=list(range(largest_order))) + + return mat + # pylint: disable=useless-super-delegation class FermiSentence(dict): @@ -463,6 +493,36 @@ def simplify(self, tol=1e-8): if abs(coeff) <= tol: del self[fw] + def to_mat(self, n_orbitals=None): + r"""Return the matrix representation. + + Args: + n_orbitals (int or None): Number of orbitals. If not provided, it will be inferred from + the largest orbital index in the Fermi operator + + Returns: + NumpyArray: Matrix representation of the :class:`~.FermiSentence`. + + **Example** + + >>> fs = FermiSentence({FermiWord({(0, 0): "+", (1, 1): "-"}): 1.2, FermiWord({(0, 0): "+", (1, 0): "-"}): 3.1}) + >>> fs.to_mat() + array([0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j], + [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j], + [0.0 + 0.0j, 1.2 + 0.0j, 3.1 + 0.0j, 0.0 + 0.0j], + [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 3.1 + 0.0j]) + """ + largest_orb_id = max(key[1] for fermi_word in self.keys() for key in fermi_word.keys()) + 1 + if n_orbitals and n_orbitals < largest_orb_id: + raise ValueError( + f"n_orbitals cannot be smaller than {largest_orb_id}, got: {n_orbitals}." + ) + + largest_order = n_orbitals or largest_orb_id + mat = qml.jordan_wigner(self, ps=True).to_mat(wire_order=list(range(largest_order))) + + return mat + def from_string(fermi_string): r"""Return a fermionic operator object from its string representation. diff --git a/pennylane/fourier/circuit_spectrum.py b/pennylane/fourier/circuit_spectrum.py index 5e695eae139..75aaa3444ff 100644 --- a/pennylane/fourier/circuit_spectrum.py +++ b/pennylane/fourier/circuit_spectrum.py @@ -15,10 +15,10 @@ of a quantum circuit, that is the frequencies without considering preprocessing in the QNode.""" from functools import partial -from typing import Callable, Sequence from pennylane import transform -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .utils import get_spectrum, join_spectra @@ -26,7 +26,7 @@ @partial(transform, is_informative=True) def circuit_spectrum( tape: QuantumTape, encoding_gates=None, decimals=8 -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Compute the frequency spectrum of the Fourier representation of simple quantum circuits ignoring classical preprocessing. diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py index 9027c869393..3ec088efef0 100644 --- a/pennylane/gradients/__init__.py +++ b/pennylane/gradients/__init__.py @@ -318,8 +318,11 @@ def circuit(weights): .. code-block:: python + from pennylane.tape import QuantumTapeBatch + from pennylane.typing import PostprocessingFn + @transform - def my_custom_gradient(tape: qml.tape.QuantumTape, **kwargs) -> (Sequence[qml.tape.QuantumTape], Callable): + def my_custom_gradient(tape: qml.tape.QuantumTape, **kwargs) -> tuple[QuantumTapeBatch, PostprocessingFn]: ... return gradient_tapes, processing_fn diff --git a/pennylane/gradients/adjoint_metric_tensor.py b/pennylane/gradients/adjoint_metric_tensor.py index 5c16528b596..687bdfc139a 100644 --- a/pennylane/gradients/adjoint_metric_tensor.py +++ b/pennylane/gradients/adjoint_metric_tensor.py @@ -16,7 +16,6 @@ """ from functools import partial from itertools import chain -from typing import Callable, Sequence import numpy as np @@ -24,7 +23,9 @@ # pylint: disable=too-many-statements,unused-argument from pennylane.gradients.metric_tensor import _contract_metric_tensor_with_cjac +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn def _reshape_real_imag(state, dim): @@ -59,7 +60,7 @@ def _group_operations(tape): def _expand_trainable_multipar( tape: qml.tape.QuantumTape, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand trainable multi-parameter operations in a quantum tape.""" interface = qml.math.get_interface(*tape.get_parameters()) @@ -77,7 +78,9 @@ def _expand_trainable_multipar( is_informative=True, use_argnum_in_expand=True, ) -def adjoint_metric_tensor(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable): +def adjoint_metric_tensor( + tape: qml.tape.QuantumTape, +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Implements the adjoint method outlined in `Jones `__ to compute the metric tensor. diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py index 05a353b44cd..bad2d650c0c 100644 --- a/pennylane/gradients/finite_difference.py +++ b/pennylane/gradients/finite_difference.py @@ -19,7 +19,6 @@ from functools import partial # pylint: disable=protected-access,too-many-arguments,too-many-branches,too-many-statements,unused-argument -from typing import Callable, Sequence from warnings import warn import numpy as np @@ -30,6 +29,8 @@ from pennylane import transform from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac from pennylane.measurements import ProbabilityMP +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .general_shift_rules import generate_shifted_tapes from .gradient_transform import ( @@ -196,7 +197,7 @@ def _expand_transform_finite_diff( strategy="forward", f0=None, validate_params=True, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand function to be applied before finite difference.""" [new_tape], postprocessing = qml.devices.preprocess.decompose( tape, @@ -227,7 +228,7 @@ def finite_diff( strategy="forward", f0=None, validate_params=True, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the finite-difference gradient of all gate parameters with respect to its inputs. Args: diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py index 223ffe1f43b..be3595a1d85 100644 --- a/pennylane/gradients/gradient_transform.py +++ b/pennylane/gradients/gradient_transform.py @@ -149,8 +149,12 @@ def _try_zero_grad_from_graph_or_get_grad_method(tape, param_index, use_graph=Tr par_info = tape.par_info[param_index] if use_graph: - op_or_mp = tape[par_info["op_idx"]] - if not any(tape.graph.has_path(op_or_mp, mp) for mp in tape.measurements): + op_or_mp_idx = par_info["op_idx"] + n_ops = len(tape.operations) + if not any( + tape.graph.has_path_idx(op_or_mp_idx, n_ops + i) + for i, _ in enumerate(tape.measurements) + ): # there is no influence of this operation on any of the observables return "0" diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py index bf70b3a7f2e..da1157400c6 100644 --- a/pennylane/gradients/hadamard_gradient.py +++ b/pennylane/gradients/hadamard_gradient.py @@ -17,16 +17,15 @@ """ from functools import partial -# pylint: disable=unused-argument -from typing import Callable, Sequence - import numpy as np import pennylane as qml from pennylane import transform from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac from pennylane.gradients.metric_tensor import _get_aux_wire +from pennylane.tape import QuantumTapeBatch from pennylane.transforms.tape_expand import expand_invalid_trainable_hadamard_gradient +from pennylane.typing import PostprocessingFn from .gradient_transform import ( _all_zero_grad, @@ -38,13 +37,15 @@ find_and_validate_gradient_methods, ) +# pylint: disable=unused-argument + def _expand_transform_hadamard( tape: qml.tape.QuantumTape, argnum=None, aux_wire=None, device_wires=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand function to be applied before hadamard gradient.""" expanded_tape = expand_invalid_trainable_hadamard_gradient(tape) @@ -68,7 +69,7 @@ def hadamard_grad( argnum=None, aux_wire=None, device_wires=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the Hadamard test gradient of all gates with respect to their inputs. diff --git a/pennylane/gradients/metric_tensor.py b/pennylane/gradients/metric_tensor.py index 5489f353c2e..5d139a1de41 100644 --- a/pennylane/gradients/metric_tensor.py +++ b/pennylane/gradients/metric_tensor.py @@ -18,14 +18,15 @@ import functools import warnings from functools import partial -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane.circuit_graph import LayerData from pennylane.queuing import WrappedObj +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn def _contract_metric_tensor_with_cjac(mt, cjac, tape): # pylint: disable=unused-argument @@ -74,7 +75,7 @@ def _expand_metric_tensor( allow_nonunitary=True, aux_wire=None, device_wires=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): # pylint: disable=too-many-arguments +) -> tuple[QuantumTapeBatch, PostprocessingFn]: # pylint: disable=too-many-arguments """Set the metric tensor based on whether non-unitary gates are allowed.""" # pylint: disable=unused-argument,too-many-arguments @@ -96,7 +97,7 @@ def metric_tensor( # pylint:disable=too-many-arguments allow_nonunitary=True, aux_wire=None, device_wires=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Returns a function that computes the metric tensor of a given QNode or quantum tape. The metric tensor convention we employ here has the following form: diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py index 63152ee3d0f..8b87b20e4bc 100644 --- a/pennylane/gradients/parameter_shift.py +++ b/pennylane/gradients/parameter_shift.py @@ -17,14 +17,13 @@ """ from functools import partial -# pylint: disable=protected-access,too-many-arguments,too-many-statements,unused-argument -from typing import Callable, Sequence - import numpy as np import pennylane as qml from pennylane import transform from pennylane.measurements import VarianceMP +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .finite_difference import finite_diff from .general_shift_rules import ( @@ -46,6 +45,9 @@ reorder_grads, ) +# pylint: disable=protected-access,too-many-arguments,too-many-statements,unused-argument + + NONINVOLUTORY_OBS = { "Hermitian": lambda obs: obs.__class__(obs.matrix() @ obs.matrix(), wires=obs.wires), "SparseHamiltonian": lambda obs: obs.__class__(obs.matrix() @ obs.matrix(), wires=obs.wires), @@ -772,7 +774,7 @@ def _expand_transform_param_shift( fallback_fn=finite_diff, f0=None, broadcast=False, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand function to be applied before parameter shift.""" [new_tape], postprocessing = qml.devices.preprocess.decompose( tape, @@ -802,7 +804,7 @@ def param_shift( fallback_fn=finite_diff, f0=None, broadcast=False, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the parameter-shift gradient of all gate parameters with respect to its inputs. diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py index 2eb69e3f168..31db666be2c 100644 --- a/pennylane/gradients/parameter_shift_cv.py +++ b/pennylane/gradients/parameter_shift_cv.py @@ -19,9 +19,6 @@ import warnings from functools import partial -# pylint: disable=protected-access,too-many-arguments,too-many-statements,too-many-branches,unused-argument -from typing import Callable, Sequence - import numpy as np import pennylane as qml @@ -38,13 +35,17 @@ StateMP, VarianceMP, ) +from pennylane.tape import QuantumTapeBatch from pennylane.transforms.tape_expand import expand_invalid_trainable +from pennylane.typing import PostprocessingFn from .finite_difference import finite_diff from .general_shift_rules import generate_shifted_tapes, process_shifts from .gradient_transform import _no_trainable_grad from .parameter_shift import _get_operation_recipe, expval_param_shift +# pylint: disable=protected-access,too-many-arguments,too-many-statements,too-many-branches,unused-argument + def _grad_method_cv(tape, idx): """Determine the best CV parameter-shift gradient recipe for a given @@ -376,8 +377,10 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient B_inv = B.copy() succ = tape.graph.descendants_in_order((op,)) - operation_descendents = itertools.filterfalse(qml.circuit_graph._is_observable, succ) - observable_descendents = filter(qml.circuit_graph._is_observable, succ) + operation_descendents = itertools.filterfalse( + lambda obj: isinstance(obj, MeasurementProcess), succ + ) + observable_descendents = filter(lambda obj: isinstance(obj, MeasurementProcess), succ) for BB in operation_descendents: if not BB.supports_heisenberg: @@ -501,7 +504,7 @@ def _expand_transform_param_shift_cv( fallback_fn=finite_diff, f0=None, force_order2=False, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand function to be applied before parameter shift CV.""" expanded_tape = expand_invalid_trainable(tape) @@ -529,7 +532,7 @@ def param_shift_cv( fallback_fn=finite_diff, f0=None, force_order2=False, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a continuous-variable QNode to compute the parameter-shift gradient of all gate parameters with respect to its inputs. diff --git a/pennylane/gradients/parameter_shift_hessian.py b/pennylane/gradients/parameter_shift_hessian.py index b875d3942cd..b07f502b104 100644 --- a/pennylane/gradients/parameter_shift_hessian.py +++ b/pennylane/gradients/parameter_shift_hessian.py @@ -19,13 +19,14 @@ import warnings from functools import partial from string import ascii_letters as ABC -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane.measurements import ProbabilityMP, StateMP, VarianceMP +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from .general_shift_rules import ( _combine_shift_rules, @@ -440,7 +441,7 @@ def _contract_qjac_with_cjac(qhess, cjac, tape): @partial(transform, classical_cotransform=_contract_qjac_with_cjac, final_transform=True) def param_shift_hessian( tape: qml.tape.QuantumTape, argnum=None, diagonal_shifts=None, off_diagonal_shifts=None, f0=None -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the parameter-shift Hessian with respect to its trainable parameters. This is the Hessian transform to replace the old one in the new return types system diff --git a/pennylane/gradients/pulse_gradient.py b/pennylane/gradients/pulse_gradient.py index 2cc0852fbb7..b3b154608c3 100644 --- a/pennylane/gradients/pulse_gradient.py +++ b/pennylane/gradients/pulse_gradient.py @@ -17,13 +17,14 @@ """ import warnings from functools import partial -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane import transform from pennylane.pulse import HardwareHamiltonian, ParametrizedEvolution +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .general_shift_rules import eigvals_to_frequencies, generate_shift_rule from .gradient_transform import ( @@ -291,7 +292,7 @@ def stoch_pulse_grad( num_split_times=1, sampler_seed=None, use_broadcasting=False, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Compute the gradient of a quantum circuit composed of pulse sequences by applying the stochastic parameter shift rule. @@ -392,7 +393,7 @@ def stoch_pulse_grad( jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit.jax") + dev = qml.device("default.qubit") def sin(p, t): return jax.numpy.sin(p * t) diff --git a/pennylane/gradients/pulse_gradient_odegen.py b/pennylane/gradients/pulse_gradient_odegen.py index 4e899b33314..d01baaadd80 100644 --- a/pennylane/gradients/pulse_gradient_odegen.py +++ b/pennylane/gradients/pulse_gradient_odegen.py @@ -16,7 +16,6 @@ parameter-shift gradient of pulse sequences in a qubit-based quantum tape. """ from functools import partial -from typing import Callable, Sequence import numpy as np @@ -24,6 +23,8 @@ from pennylane import transform from pennylane.ops.qubit.special_unitary import _pauli_decompose, pauli_basis_strings from pennylane.pulse import ParametrizedEvolution +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .gradient_transform import ( _all_zero_grad, @@ -402,7 +403,7 @@ def processing_fn(results): @partial(transform, final_transform=True) def pulse_odegen( tape: qml.tape.QuantumTape, argnum=None, atol=1e-7 -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the pulse generator parameter-shift gradient of pulses in a pulse program with respect to their inputs. This method combines automatic differentiation of few-qubit operations with @@ -484,7 +485,7 @@ def pulse_odegen( .. code-block:: python - dev = qml.device("default.qubit.jax") + dev = qml.device("default.qubit") @qml.qnode(dev, interface="jax", diff_method=qml.gradients.pulse_odegen) def circuit(params): diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py index a90cc062eb8..4a550007590 100644 --- a/pennylane/gradients/spsa_gradient.py +++ b/pennylane/gradients/spsa_gradient.py @@ -17,15 +17,14 @@ """ from functools import partial -# pylint: disable=protected-access,too-many-arguments,too-many-branches,too-many-statements,unused-argument -from typing import Callable, Sequence - import numpy as np import pennylane as qml from pennylane import transform from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac +from pennylane.tape import QuantumTapeBatch from pennylane.transforms.tape_expand import expand_invalid_trainable +from pennylane.typing import PostprocessingFn from .finite_difference import _processing_fn, finite_diff_coeffs from .general_shift_rules import generate_multishifted_tapes @@ -37,6 +36,8 @@ find_and_validate_gradient_methods, ) +# pylint: disable=protected-access,too-many-arguments,too-many-branches,too-many-statements,unused-argument + def _rademacher_sampler(indices, num_params, *args, rng): r"""Sample a random vector with (independent) entries from {+1, -1} with balanced probability. @@ -72,7 +73,7 @@ def _expand_transform_spsa( num_directions=1, sampler=_rademacher_sampler, sampler_rng=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Expand function to be applied before spsa gradient.""" expanded_tape = expand_invalid_trainable(tape) @@ -103,7 +104,7 @@ def spsa_grad( num_directions=1, sampler=_rademacher_sampler, sampler_rng=None, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Transform a circuit to compute the SPSA gradient of all gate parameters with respect to its inputs. This estimator shifts all parameters simultaneously and approximates the gradient based on these shifts and a diff --git a/pennylane/logging/formatters/formatter.py b/pennylane/logging/formatters/formatter.py index 09daa5e9570..75fc8c42559 100644 --- a/pennylane/logging/formatters/formatter.py +++ b/pennylane/logging/formatters/formatter.py @@ -15,7 +15,7 @@ import inspect import logging from logging import Formatter -from typing import NamedTuple, Tuple, Union +from typing import NamedTuple, Optional # For color-code definitions see https://en.wikipedia.org/wiki/ANSI_escape_code#Colors # 24-bit mode support RGB color codes in 8bit-wide r;g;b format @@ -38,14 +38,14 @@ class ColorScheme(NamedTuple): warning: str error: str critical: str - debug_bg: Union[None, str] - info_bg: Union[None, str] - warning_bg: Union[None, str] - error_bg: Union[None, str] - critical_bg: Union[None, str] + debug_bg: Optional[str] + info_bg: Optional[str] + warning_bg: Optional[str] + error_bg: Optional[str] + critical_bg: Optional[str] -def build_code_rgb(rgb: Tuple[int, int, int], rgb_bg: Union[None, Tuple[int, int, int]] = None): +def build_code_rgb(rgb: tuple[int, int, int], rgb_bg: Optional[tuple[int, int, int]] = None): """ Utility function to generate the appropriate ANSI RGB codes for a given set of foreground (font) and background colors. """ diff --git a/pennylane/math/matrix_manipulation.py b/pennylane/math/matrix_manipulation.py index 99a2ec811d5..fb06140011d 100644 --- a/pennylane/math/matrix_manipulation.py +++ b/pennylane/math/matrix_manipulation.py @@ -14,8 +14,8 @@ """This module contains methods to expand the matrix representation of an operator to a higher hilbert space with re-ordered wires.""" import itertools +from collections.abc import Callable, Generator, Iterable from functools import reduce -from typing import Generator, Iterable, Tuple import numpy as np from scipy.sparse import csr_matrix, eye, kron @@ -280,8 +280,8 @@ def _permutation_sparse_matrix(expanded_wires: Iterable, wire_order: Iterable) - def reduce_matrices( - mats_and_wires_gen: Generator[Tuple[np.ndarray, Wires], None, None], reduce_func: callable -) -> Tuple[np.ndarray, Wires]: + mats_and_wires_gen: Generator[tuple[np.ndarray, Wires], None, None], reduce_func: Callable +) -> tuple[np.ndarray, Wires]: """Apply the given ``reduce_func`` cumulatively to the items of the ``mats_and_wires_gen`` generator, from left to right, so as to reduce the sequence to a tuple containing a single matrix and the wires it acts on. @@ -295,7 +295,7 @@ def reduce_matrices( Tuple[tensor, Wires]: a tuple containing the reduced matrix and the wires it acts on """ - def expand_and_reduce(op1_tuple: Tuple[np.ndarray, Wires], op2_tuple: Tuple[np.ndarray, Wires]): + def expand_and_reduce(op1_tuple: tuple[np.ndarray, Wires], op2_tuple: tuple[np.ndarray, Wires]): mat1, wires1 = op1_tuple mat2, wires2 = op2_tuple expanded_wires = wires1 + wires2 diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py index 08a4d5d40df..7ac08e7a758 100644 --- a/pennylane/measurements/classical_shadow.py +++ b/pennylane/measurements/classical_shadow.py @@ -15,9 +15,9 @@ This module contains the qml.classical_shadow measurement. """ import copy -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from string import ascii_letters as ABC -from typing import Optional, Sequence, Union +from typing import Optional, Union import numpy as np diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py index 331850e0e61..0a1669d7403 100644 --- a/pennylane/measurements/counts.py +++ b/pennylane/measurements/counts.py @@ -14,7 +14,8 @@ """ This module contains the qml.counts measurement. """ -from typing import Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Optional import numpy as np @@ -238,8 +239,8 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, - bin_size: int = None, + shot_range: Optional[tuple[int, ...]] = None, + bin_size: Optional[int] = None, ): with qml.queuing.QueuingManager.stop_recording(): samples = qml.sample(op=self.obs or self.mv, wires=self._wires).process_samples( diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py index 9b98a887ccc..95499fe9ea8 100644 --- a/pennylane/measurements/expval.py +++ b/pennylane/measurements/expval.py @@ -14,7 +14,8 @@ """ This module contains the qml.expval measurement. """ -from typing import Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Optional, Union import pennylane as qml from pennylane.operation import Operator @@ -107,8 +108,8 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, - bin_size: int = None, + shot_range: Optional[tuple[int, ...]] = None, + bin_size: Optional[int] = None, ): if not self.wires: return qml.math.squeeze(self.eigvals()) diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py index ccf87e59c5f..b2401492f73 100644 --- a/pennylane/measurements/measurements.py +++ b/pennylane/measurements/measurements.py @@ -19,8 +19,9 @@ import copy import functools from abc import ABC, abstractmethod +from collections.abc import Sequence from enum import Enum -from typing import Optional, Sequence, Tuple, Union +from typing import Optional, Union import pennylane as qml from pennylane.math.utils import is_abstract @@ -288,7 +289,7 @@ def numeric_type(self) -> type: f"The numeric type of the measurement {self.__class__.__name__} is not defined." ) - def shape(self, device, shots: Shots) -> Tuple: + def shape(self, device, shots: Shots) -> tuple: """The expected output shape of the MeasurementProcess. Note that the output shape is dependent on the shots or device when: @@ -604,7 +605,7 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, + shot_range: tuple[int] = None, bin_size: int = None, ): """Process the given samples. diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py index 8e51e130d3f..95c71f19615 100644 --- a/pennylane/measurements/mutual_info.py +++ b/pennylane/measurements/mutual_info.py @@ -15,8 +15,9 @@ """ This module contains the qml.mutual_info measurement. """ +from collections.abc import Sequence from copy import copy -from typing import Optional, Sequence +from typing import Optional import pennylane as qml from pennylane.wires import Wires diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py index fdc5bf4aca9..74e62bf122d 100644 --- a/pennylane/measurements/probs.py +++ b/pennylane/measurements/probs.py @@ -14,7 +14,8 @@ """ This module contains the qml.probs measurement. """ -from typing import Sequence, Tuple +from collections.abc import Sequence +from typing import Optional import numpy as np @@ -177,8 +178,8 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, - bin_size: int = None, + shot_range: Optional[tuple[int, ...]] = None, + bin_size: Optional[int] = None, ): wire_map = dict(zip(wire_order, range(len(wire_order)))) mapped_wires = [wire_map[w] for w in self.wires] diff --git a/pennylane/measurements/purity.py b/pennylane/measurements/purity.py index 281091a3033..e071f05dd1b 100644 --- a/pennylane/measurements/purity.py +++ b/pennylane/measurements/purity.py @@ -15,8 +15,8 @@ """ This module contains the qml.purity measurement. """ - -from typing import Optional, Sequence +from collections.abc import Sequence +from typing import Optional import pennylane as qml from pennylane.wires import Wires diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py index f403230d433..b50c7f0a08d 100644 --- a/pennylane/measurements/sample.py +++ b/pennylane/measurements/sample.py @@ -15,7 +15,8 @@ This module contains the qml.sample measurement. """ import functools -from typing import Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Optional, Union import numpy as np @@ -254,7 +255,7 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, + shot_range: tuple[int, ...] = None, bin_size: int = None, ): wire_map = dict(zip(wire_order, range(len(wire_order)))) diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py index 2cbaea69554..b6993427179 100644 --- a/pennylane/measurements/state.py +++ b/pennylane/measurements/state.py @@ -14,7 +14,8 @@ """ This module contains the qml.state measurement. """ -from typing import Optional, Sequence +from collections.abc import Sequence +from typing import Optional import pennylane as qml from pennylane.typing import TensorLike diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py index 19aa167c1ca..21c9bf1345f 100644 --- a/pennylane/measurements/var.py +++ b/pennylane/measurements/var.py @@ -15,7 +15,8 @@ """ This module contains the qml.var measurement. """ -from typing import Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Optional, Union import pennylane as qml from pennylane.operation import Operator @@ -99,8 +100,8 @@ def process_samples( self, samples: Sequence[complex], wire_order: Wires, - shot_range: Tuple[int] = None, - bin_size: int = None, + shot_range: Optional[tuple[int, ...]] = None, + bin_size: Optional[int] = None, ): # estimate the variance op = self.mv if self.mv is not None else self.obs diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py index 52c7b9a1734..1589cee7b90 100644 --- a/pennylane/measurements/vn_entropy.py +++ b/pennylane/measurements/vn_entropy.py @@ -15,7 +15,8 @@ """ This module contains the qml.vn_entropy measurement. """ -from typing import Optional, Sequence +from collections.abc import Sequence +from typing import Optional import pennylane as qml from pennylane.wires import Wires diff --git a/pennylane/operation.py b/pennylane/operation.py index d31933747cb..00d2ce9dde0 100644 --- a/pennylane/operation.py +++ b/pennylane/operation.py @@ -249,7 +249,7 @@ import warnings from contextlib import contextmanager from enum import IntEnum -from typing import List, Optional, Tuple +from typing import Optional import numpy as np from numpy.linalg import multi_dot @@ -1089,7 +1089,6 @@ def __init__(self, *params, wires=None, id=None): self._name = self.__class__.__name__ #: str: name of the operator self._id = id - self.queue_idx = None #: int, None: index of the Operator in the circuit queue, or None if not in a queue self._pauli_rep = None # Union[PauliSentence, None]: Representation of the operator as a pauli sentence, if applicable wires_from_args = False @@ -1305,7 +1304,7 @@ def has_decomposition(cls): or cls.decomposition != Operator.decomposition ) - def decomposition(self) -> List["Operator"]: + def decomposition(self) -> list["Operator"]: r"""Representation of the operator as a product of other operators. .. math:: O = O_1 O_2 \dots O_n @@ -1322,7 +1321,7 @@ def decomposition(self) -> List["Operator"]: ) @staticmethod - def compute_decomposition(*params, wires=None, **hyperparameters) -> List["Operator"]: + def compute_decomposition(*params, wires=None, **hyperparameters) -> list["Operator"]: r"""Representation of the operator as a product of other operators (static method). .. math:: O = O_1 O_2 \dots O_n. @@ -1361,7 +1360,7 @@ def has_diagonalizing_gates(cls): @staticmethod def compute_diagonalizing_gates( *params, wires, **hyperparams - ) -> List["Operator"]: # pylint: disable=unused-argument + ) -> list["Operator"]: # pylint: disable=unused-argument r"""Sequence of gates that diagonalize the operator in the computational basis (static method). Given the eigendecomposition :math:`O = U \Sigma U^{\dagger}` where @@ -1435,7 +1434,7 @@ def generator(self): # pylint: disable=no-self-use """ raise GeneratorUndefinedError(f"Operation {self.name} does not have a generator") - def pow(self, z) -> List["Operator"]: + def pow(self, z) -> list["Operator"]: """A list of new operators equal to this one raised to the given power. Args: @@ -1499,9 +1498,19 @@ def adjoint(self) -> "Operator": # pylint:disable=no-self-use def expand(self): """Returns a tape that contains the decomposition of the operator. + .. warning:: + This function is deprecated and will be removed in version 0.39. + The same behaviour can be achieved simply through 'qml.tape.QuantumScript(self.decomposition())'. + Returns: .QuantumTape: quantum tape """ + warnings.warn( + "'Operator.expand' is deprecated and will be removed in version 0.39. " + "The same behaviour can be achieved simply through 'qml.tape.QuantumScript(self.decomposition())'.", + qml.PennyLaneDeprecationWarning, + ) + if not self.has_decomposition: raise DecompositionUndefinedError @@ -1760,7 +1769,7 @@ def control_wires(self): # pragma: no cover """ return Wires([]) - def single_qubit_rot_angles(self) -> Tuple[float, float, float]: + def single_qubit_rot_angles(self) -> tuple[float, float, float]: r"""The parameters required to implement a single-qubit gate as an equivalent ``Rot`` gate, up to a global phase. @@ -1857,7 +1866,7 @@ class Channel(Operation, abc.ABC): @abc.abstractmethod def compute_kraus_matrices( *params, **hyperparams - ) -> List[np.ndarray]: # pylint:disable=unused-argument + ) -> list[np.ndarray]: # pylint:disable=unused-argument """Kraus matrices representing a quantum channel, specified in the computational basis. @@ -2092,7 +2101,7 @@ def _primitive_bind_call(cls, *args, **kwargs): def __init__(self, *args): # pylint: disable=super-init-not-called self._eigvals_cache = None - self.obs: List[Observable] = [] + self.obs: list[Observable] = [] self._args = args self._batch_size = None self._pauli_rep = None diff --git a/pennylane/ops/functions/assert_valid.py b/pennylane/ops/functions/assert_valid.py index 860ed3d1678..bc92514a4ce 100644 --- a/pennylane/ops/functions/assert_valid.py +++ b/pennylane/ops/functions/assert_valid.py @@ -53,16 +53,13 @@ def _check_decomposition(op, skip_wire_mapping): with qml.queuing.AnnotatedQueue() as queued_decomp: op.decomposition() processed_queue = qml.tape.QuantumTape.from_queue(queued_decomp) - expand = op.expand() assert isinstance(decomp, list), "decomposition must be a list" assert isinstance(compute_decomp, list), "decomposition must be a list" - assert isinstance(expand, qml.tape.QuantumScript), "expand must return a QuantumScript" - for o1, o2, o3, o4 in zip(decomp, compute_decomp, processed_queue, expand): + for o1, o2, o3 in zip(decomp, compute_decomp, processed_queue): assert o1 == o2, "decomposition must match compute_decomposition" assert o1 == o3, "decomposition must match queued operations" - assert o1 == o4, "decomposition must match expansion" assert isinstance(o1, qml.operation.Operator), "decomposition must contain operators" if skip_wire_mapping: @@ -83,9 +80,6 @@ def _check_decomposition(op, skip_wire_mapping): qml.operation.DecompositionUndefinedError, failure_comment=failure_comment, )() - _assert_error_raised( - op.expand, qml.operation.DecompositionUndefinedError, failure_comment=failure_comment - )() _assert_error_raised( op.compute_decomposition, qml.operation.DecompositionUndefinedError, diff --git a/pennylane/ops/functions/bind_new_parameters.py b/pennylane/ops/functions/bind_new_parameters.py index c80456ab559..af29d1576af 100644 --- a/pennylane/ops/functions/bind_new_parameters.py +++ b/pennylane/ops/functions/bind_new_parameters.py @@ -17,8 +17,9 @@ # pylint: disable=missing-docstring import copy +from collections.abc import Sequence from functools import singledispatch -from typing import Sequence, Union +from typing import Union import pennylane as qml from pennylane.operation import Operator, Tensor diff --git a/pennylane/ops/functions/dot.py b/pennylane/ops/functions/dot.py index b6bed90c72e..a49dd668c10 100644 --- a/pennylane/ops/functions/dot.py +++ b/pennylane/ops/functions/dot.py @@ -17,7 +17,8 @@ """ # pylint: disable=too-many-branches from collections import defaultdict -from typing import Callable, Sequence, Union +from collections.abc import Callable, Sequence +from typing import Union import pennylane as qml from pennylane.operation import Operator, convert_to_opmath diff --git a/pennylane/ops/functions/eigvals.py b/pennylane/ops/functions/eigvals.py index 6e64ccc3d69..ff6d0224e74 100644 --- a/pennylane/ops/functions/eigvals.py +++ b/pennylane/ops/functions/eigvals.py @@ -18,14 +18,14 @@ # pylint: disable=protected-access from functools import partial, reduce -from typing import Callable, Sequence import scipy import pennylane as qml from pennylane import transform +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import TransformError -from pennylane.typing import TensorLike +from pennylane.typing import PostprocessingFn, TensorLike def eigvals(op: qml.operation.Operator, k=1, which="SA") -> TensorLike: @@ -135,13 +135,13 @@ def circuit(theta): try: return op.eigvals() except qml.operation.EigvalsUndefinedError: - return eigvals(op.expand(), k=k, which=which) + return eigvals(qml.tape.QuantumScript(op.decomposition()), k=k, which=which) @partial(transform, is_informative=True) def _eigvals_tranform( tape: qml.tape.QuantumTape, k=1, which="SA" -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: def processing_fn(res): [qs] = res op_wires = [op.wires for op in qs.operations] diff --git a/pennylane/ops/functions/evolve.py b/pennylane/ops/functions/evolve.py index c7d2722b9b2..96660aea77e 100644 --- a/pennylane/ops/functions/evolve.py +++ b/pennylane/ops/functions/evolve.py @@ -129,7 +129,10 @@ def evolve(*args, **kwargs): # pylint: disable=unused-argument import jax - dev = qml.device("default.qubit.jax", wires=4) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit") + @jax.jit @qml.qnode(dev, interface="jax") def circuit(params): @@ -138,13 +141,13 @@ def circuit(params): >>> params = [1., 2., 3., 4.] >>> circuit(params) - Array(0.8627419, dtype=float32) + Array(0.86231063, dtype=float64) >>> jax.grad(circuit)(params) - [Array(50.690746, dtype=float32), - Array(-6.296886e-05, dtype=float32), - Array(-6.3341584e-05, dtype=float32), - Array(-7.052516e-05, dtype=float32)] + [Array(50.391273, dtype=float64), + Array(-9.42415807e-05, dtype=float64), + Array(-0.0001049, dtype=float64), + Array(-0.00010601, dtype=float64)] .. note:: In the example above, the decorator ``@jax.jit`` is used to compile this execution just-in-time. This means diff --git a/pennylane/ops/functions/map_wires.py b/pennylane/ops/functions/map_wires.py index 8839e94a36f..162b0a586b4 100644 --- a/pennylane/ops/functions/map_wires.py +++ b/pennylane/ops/functions/map_wires.py @@ -14,15 +14,17 @@ """ This module contains the qml.map_wires function. """ +from collections.abc import Callable from functools import partial -from typing import Callable, Sequence, Union +from typing import Union import pennylane as qml from pennylane import transform from pennylane.measurements import MeasurementProcess from pennylane.operation import Operator from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn from pennylane.workflow import QNode @@ -108,7 +110,7 @@ def map_wires( @partial(transform) def _map_wires_transform( tape: qml.tape.QuantumTape, wire_map=None, queue=False -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: ops = [ ( map_wires(op, wire_map, queue=queue) diff --git a/pennylane/ops/functions/matrix.py b/pennylane/ops/functions/matrix.py index eafb8dc38a1..66758558a00 100644 --- a/pennylane/ops/functions/matrix.py +++ b/pennylane/ops/functions/matrix.py @@ -17,14 +17,15 @@ from functools import partial # pylint: disable=protected-access,too-many-branches -from typing import Callable, Sequence, Union +from typing import Union import pennylane as qml from pennylane import transform from pennylane.operation import Operator from pennylane.pauli import PauliSentence, PauliWord +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import TransformError -from pennylane.typing import TensorLike +from pennylane.typing import PostprocessingFn, TensorLike def catalyst_qjit(qnode): @@ -227,13 +228,13 @@ def circuit(): try: return op.matrix(wire_order=wire_order) except: # pylint: disable=bare-except - return matrix(op.expand(), wire_order=wire_order or op.wires) + return matrix(qml.tape.QuantumScript(op.decomposition()), wire_order=wire_order or op.wires) @partial(transform, is_informative=True) def _matrix_transform( tape: qml.tape.QuantumTape, wire_order=None, **kwargs -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: if not tape.wires: raise qml.operation.MatrixUndefinedError diff --git a/pennylane/ops/functions/simplify.py b/pennylane/ops/functions/simplify.py index e7d7c830517..52a4e570796 100644 --- a/pennylane/ops/functions/simplify.py +++ b/pennylane/ops/functions/simplify.py @@ -14,14 +14,16 @@ """ This module contains the qml.simplify function. """ +from collections.abc import Callable from copy import copy -from typing import Callable, Sequence, Union +from typing import Union import pennylane as qml from pennylane.measurements import MeasurementProcess from pennylane.operation import Operator from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn from pennylane.workflow import QNode @@ -98,7 +100,7 @@ def simplify(input: Union[Operator, MeasurementProcess, QuantumTape, QNode, Call @qml.transform -def _simplify_transform(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def _simplify_transform(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: with qml.QueuingManager.stop_recording(): new_operations = [op.simplify() for op in tape.operations] new_measurements = [m.simplify() for m in tape.measurements] diff --git a/pennylane/ops/op_math/adjoint.py b/pennylane/ops/op_math/adjoint.py index 0ccd25113be..d7f4d220ead 100644 --- a/pennylane/ops/op_math/adjoint.py +++ b/pennylane/ops/op_math/adjoint.py @@ -85,7 +85,7 @@ def adjoint(fn, lazy=True): ... return qml.expval(qml.Z(0)) >>> print(qml.draw(circuit2)("y")) 0: ──RY(y)†── - >>> print(qml.draw(circuit2, expansion_strategy="device")(0.1)) + >>> print(qml.draw(circuit2, level="device")(0.1)) 0: ──RY(-0.10)── The adjoint transforms can also be used to apply the adjoint of diff --git a/pennylane/ops/op_math/composite.py b/pennylane/ops/op_math/composite.py index 38636be6ea9..9bada443a57 100644 --- a/pennylane/ops/op_math/composite.py +++ b/pennylane/ops/op_math/composite.py @@ -17,7 +17,7 @@ # pylint: disable=too-many-instance-attributes,invalid-sequence-index import abc import copy -from typing import Callable, List +from collections.abc import Callable import pennylane as qml from pennylane import math @@ -59,7 +59,6 @@ def __init__( self, *operands: Operator, id=None, _pauli_rep=None ): # pylint: disable=super-init-not-called self._id = id - self.queue_idx = None self._name = self.__class__.__name__ self.operands = operands @@ -189,7 +188,7 @@ def matrix(self, wire_order=None): """Representation of the operator as a matrix in the computational basis.""" @property - def overlapping_ops(self) -> List[List[Operator]]: + def overlapping_ops(self) -> list[list[Operator]]: """Groups all operands of the composite operator that act on overlapping wires. Returns: @@ -340,7 +339,7 @@ def queue(self, context=qml.QueuingManager): @classmethod @abc.abstractmethod - def _sort(cls, op_list, wire_map: dict = None) -> List[Operator]: + def _sort(cls, op_list, wire_map: dict = None) -> list[Operator]: """Sort composite operands by their wire indices.""" @property diff --git a/pennylane/ops/op_math/controlled.py b/pennylane/ops/op_math/controlled.py index d6c35e7dcdf..b065d99df74 100644 --- a/pennylane/ops/op_math/controlled.py +++ b/pennylane/ops/op_math/controlled.py @@ -19,7 +19,6 @@ from copy import copy from functools import wraps from inspect import signature -from typing import List import numpy as np from scipy import sparse @@ -750,7 +749,7 @@ def _decompose_pauli_x_based_no_control_values(op: Controlled): ) -def _decompose_custom_ops(op: Controlled) -> List["operation.Operator"]: +def _decompose_custom_ops(op: Controlled) -> list["operation.Operator"]: """Custom handling for decomposing a controlled operation""" pauli_x_based_ctrl_ops = _get_pauli_x_based_ops() @@ -789,7 +788,7 @@ def _decompose_custom_ops(op: Controlled) -> List["operation.Operator"]: return None -def _decompose_no_control_values(op: Controlled) -> List["operation.Operator"]: +def _decompose_no_control_values(op: Controlled) -> list["operation.Operator"]: """Decompose without considering control values. Returns None if no decomposition.""" decomp = _decompose_custom_ops(op) diff --git a/pennylane/ops/op_math/controlled_ops.py b/pennylane/ops/op_math/controlled_ops.py index c3a05208751..d2c5d3afb4a 100644 --- a/pennylane/ops/op_math/controlled_ops.py +++ b/pennylane/ops/op_math/controlled_ops.py @@ -16,8 +16,8 @@ """ # pylint: disable=no-value-for-parameter, arguments-differ, arguments-renamed import warnings +from collections.abc import Iterable from functools import lru_cache -from typing import Iterable import numpy as np from scipy.linalg import block_diag diff --git a/pennylane/ops/op_math/exp.py b/pennylane/ops/op_math/exp.py index 5662199c2eb..07b2cab7066 100644 --- a/pennylane/ops/op_math/exp.py +++ b/pennylane/ops/op_math/exp.py @@ -14,7 +14,6 @@ """ This submodule defines the symbolic operation that stands for an exponential of an operator. """ -from typing import List from warnings import warn import numpy as np @@ -356,7 +355,7 @@ def _pauli_rot_decomposition(base: Operator, coeff: complex): return [] return [qml.PauliRot(theta=coeff, pauli_word=pauli_word, wires=base.wires)] - def _trotter_decomposition(self, ops: List[Operator], coeffs: List[complex]): + def _trotter_decomposition(self, ops: list[Operator], coeffs: list[complex]): """Uses the Suzuki-Trotter approximation to decompose the exponential of the linear combination of ``coeffs`` and ``ops``. diff --git a/pennylane/ops/op_math/linear_combination.py b/pennylane/ops/op_math/linear_combination.py index 7fe4b69151f..8e19af77dd4 100644 --- a/pennylane/ops/op_math/linear_combination.py +++ b/pennylane/ops/op_math/linear_combination.py @@ -20,7 +20,6 @@ # pylint: disable=too-many-arguments, protected-access, too-many-instance-attributes import warnings from copy import copy -from typing import List import pennylane as qml from pennylane.operation import Observable, Operator, Tensor, convert_to_opmath @@ -119,7 +118,7 @@ def _primitive_bind_call(cls, coeffs, observables, _pauli_rep=None, **kwargs): def __init__( self, coeffs, - observables: List[Operator], + observables: list[Operator], simplify=False, grouping_type=None, method="rlf", diff --git a/pennylane/ops/op_math/prod.py b/pennylane/ops/op_math/prod.py index bf0026fd378..1bbe30e29f2 100644 --- a/pennylane/ops/op_math/prod.py +++ b/pennylane/ops/op_math/prod.py @@ -20,7 +20,7 @@ from copy import copy from functools import reduce, wraps from itertools import combinations -from typing import List, Tuple, Union +from typing import Union from scipy.sparse import kron as sparse_kron @@ -277,8 +277,8 @@ def matrix(self, wire_order=None): if self.pauli_rep: return self.pauli_rep.to_mat(wire_order=wire_order or self.wires) - mats: List[TensorLike] = [] - batched: List[bool] = [] # batched[i] tells if mats[i] is batched or not + mats: list[TensorLike] = [] + batched: list[bool] = [] # batched[i] tells if mats[i] is batched or not for ops in self.overlapping_ops: gen = ( ( @@ -358,7 +358,7 @@ def _build_pauli_rep(self): return reduce(lambda a, b: a @ b, operand_pauli_reps) return None - def _simplify_factors(self, factors: Tuple[Operator]) -> Tuple[complex, Operator]: + def _simplify_factors(self, factors: tuple[Operator]) -> tuple[complex, Operator]: """Reduces the depth of nested factors and groups identical factors. Returns: @@ -400,7 +400,7 @@ def simplify(self) -> Union["Prod", Sum]: return op if global_phase == 1 else qml.s_prod(global_phase, op).simplify() @classmethod - def _sort(cls, op_list, wire_map: dict = None) -> List[Operator]: + def _sort(cls, op_list, wire_map: dict = None) -> list[Operator]: """Insertion sort algorithm that sorts a list of product factors by their wire indices, taking into account the operator commutivity. @@ -594,7 +594,7 @@ def add(self, factor: Operator): self._add_non_pauli_factor(factor=factor, wires=wires) self._remove_pauli_factors(wires=wires) - def _add_pauli_factor(self, factor: Operator, wires: List[int]): + def _add_pauli_factor(self, factor: Operator, wires: list[int]): """Adds the given Pauli operator to the temporary ``self._pauli_factors`` dictionary. If there was another Pauli operator acting on the same wire, the two operators are grouped together using the ``self._pauli_mult`` dictionary. @@ -610,7 +610,7 @@ def _add_pauli_factor(self, factor: Operator, wires: List[int]): coeff, new_word = self._pauli_mult[old_word][op2_name] self._pauli_factors[wire] = old_coeff * coeff, new_word - def _add_non_pauli_factor(self, factor: Operator, wires: List[int]): + def _add_non_pauli_factor(self, factor: Operator, wires: list[int]): """Adds the given non-Pauli factor to the temporary ``self._non_pauli_factors`` dictionary. If there alerady exists an identical operator in the dictionary, the two are grouped together. @@ -642,7 +642,7 @@ def _add_non_pauli_factor(self, factor: Operator, wires: List[int]): self._remove_non_pauli_factors(wires=wires) self._non_pauli_factors[wires] = [op_hash, copy(exponent), factor] - def _remove_non_pauli_factors(self, wires: List[int]): + def _remove_non_pauli_factors(self, wires: list[int]): """Remove all factors from the ``self._non_pauli_factors`` dictionary that act on the given wires and add them to the ``self._factors`` tuple. @@ -662,7 +662,7 @@ def _remove_non_pauli_factors(self, wires: List[int]): if not isinstance(op, qml.Identity): self._factors += ((op,),) - def _remove_pauli_factors(self, wires: List[int]): + def _remove_pauli_factors(self, wires: list[int]): """Remove all Pauli factors from the ``self._pauli_factors`` dictionary that act on the given wires and add them to the ``self._factors`` tuple. @@ -678,7 +678,7 @@ def _remove_pauli_factors(self, wires: List[int]): self._factors += ((pauli_op,),) self.global_phase *= pauli_coeff - def remove_factors(self, wires: List[int]): + def remove_factors(self, wires: list[int]): """Remove all factors from the ``self._pauli_factors`` and ``self._non_pauli_factors`` dictionaries that act on the given wires and add them to the ``self._factors`` tuple. diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py index a8f1a0e0d29..62c705e82db 100644 --- a/pennylane/ops/op_math/sum.py +++ b/pennylane/ops/op_math/sum.py @@ -22,7 +22,6 @@ from collections import Counter from collections.abc import Iterable from copy import copy -from typing import List import pennylane as qml from pennylane import math @@ -381,7 +380,7 @@ def _build_pauli_rep(self): return None @classmethod - def _simplify_summands(cls, summands: List[Operator]): + def _simplify_summands(cls, summands: list[Operator]): """Reduces the depth of nested summands and groups equal terms together. Args: @@ -547,7 +546,7 @@ def ops(self): return ops @classmethod - def _sort(cls, op_list, wire_map: dict = None) -> List[Operator]: + def _sort(cls, op_list, wire_map: dict = None) -> list[Operator]: """Sort algorithm that sorts a list of sum summands by their wire indices. Args: diff --git a/pennylane/ops/op_math/symbolicop.py b/pennylane/ops/op_math/symbolicop.py index c1afeee1300..3053e58c8c8 100644 --- a/pennylane/ops/op_math/symbolicop.py +++ b/pennylane/ops/op_math/symbolicop.py @@ -72,7 +72,6 @@ def __copy__(self): def __init__(self, base, id=None): self.hyperparameters["base"] = base self._id = id - self.queue_idx = None self._pauli_rep = None self.queue() diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py index 521ba29584c..062a1f413de 100644 --- a/pennylane/ops/qubit/hamiltonian.py +++ b/pennylane/ops/qubit/hamiltonian.py @@ -22,7 +22,6 @@ import numbers from collections.abc import Iterable from copy import copy -from typing import List from warnings import warn import numpy as np @@ -246,7 +245,7 @@ def _primitive_bind_call(cls, coeffs, observables, **kwargs): def __init__( self, coeffs, - observables: List[Observable], + observables: list[Observable], simplify=False, grouping_type=None, _grouping_indices=None, diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py index b20e96dd478..9fd4ee43e65 100644 --- a/pennylane/ops/qubit/observables.py +++ b/pennylane/ops/qubit/observables.py @@ -584,6 +584,24 @@ def compute_diagonalizing_gates( """ return [] + @staticmethod + def compute_sparse_matrix(basis_state): # pylint: disable=arguments-differ,unused-argument + """ + Computes the sparse CSR matrix representation of the projector onto the basis state. + + Args: + basis_state (Iterable): The basis state as an iterable of integers (0 or 1). + + Returns: + scipy.sparse.csr_matrix: The sparse CSR matrix representation of the projector. + """ + + num_qubits = len(basis_state) + data = [1] + rows = [int("".join(str(bit) for bit in basis_state), 2)] + cols = rows + return csr_matrix((data, (rows, cols)), shape=(2**num_qubits, 2**num_qubits)) + class StateVectorProjector(Projector): r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where diff --git a/pennylane/optimize/adaptive.py b/pennylane/optimize/adaptive.py index c2a4502b6e5..6499b57fefb 100644 --- a/pennylane/optimize/adaptive.py +++ b/pennylane/optimize/adaptive.py @@ -13,17 +13,17 @@ # limitations under the License. """Adaptive optimizer""" import copy -from typing import Callable, Sequence # pylint: disable= no-value-for-parameter, protected-access, not-callable import pennylane as qml from pennylane import numpy as pnp from pennylane import transform -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn @transform -def append_gate(tape: QuantumTape, params, gates) -> (Sequence[QuantumTape], Callable): +def append_gate(tape: QuantumTape, params, gates) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Append parameterized gates to an existing tape. Args: diff --git a/pennylane/optimize/riemannian_gradient.py b/pennylane/optimize/riemannian_gradient.py index 063f66566ba..bf7d137cff4 100644 --- a/pennylane/optimize/riemannian_gradient.py +++ b/pennylane/optimize/riemannian_gradient.py @@ -13,7 +13,6 @@ # limitations under the License. """Riemannian gradient optimizer""" import warnings -from typing import Callable, Sequence import numpy as np from scipy.sparse.linalg import expm @@ -21,13 +20,14 @@ import pennylane as qml from pennylane import transform from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn @transform def append_time_evolution( tape: QuantumTape, riemannian_gradient, t, n, exact=False -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Append an approximate time evolution, corresponding to a Riemannian gradient on the Lie group, to an existing circuit. diff --git a/pennylane/optimize/spsa.py b/pennylane/optimize/spsa.py index e44784f7245..535f6f64675 100644 --- a/pennylane/optimize/spsa.py +++ b/pennylane/optimize/spsa.py @@ -194,6 +194,7 @@ def step_and_cost(self, objective_fn, *args, **kwargs): objective function output prior to the step. """ g = self.compute_grad(objective_fn, args, kwargs) + new_args = self.apply_grad(g, args) self.k += 1 @@ -270,7 +271,8 @@ def compute_grad(self, objective_fn, args, kwargs): shots = Shots(objective_fn.device._raw_shot_sequence) # pragma: no cover else: shots = Shots(None) - if np.prod(objective_fn.func(*args).shape(objective_fn.device, shots)) > 1: + + if np.prod(objective_fn.func(*args, **kwargs).shape(objective_fn.device, shots)) > 1: raise ValueError( "The objective function must be a scalar function for the gradient " "to be computed." diff --git a/pennylane/pauli/conversion.py b/pennylane/pauli/conversion.py index 4c4b4430f15..9e1ddbd38be 100644 --- a/pennylane/pauli/conversion.py +++ b/pennylane/pauli/conversion.py @@ -17,7 +17,7 @@ from functools import reduce, singledispatch from itertools import product from operator import matmul -from typing import Tuple, Union +from typing import Union import pennylane as qml from pennylane.math.utils import is_abstract @@ -42,7 +42,7 @@ # pylint: disable=too-many-branches def _generalized_pauli_decompose( matrix, hide_identity=False, wire_order=None, pauli=False, padding=False -) -> Tuple[qml.typing.TensorLike, list]: +) -> tuple[qml.typing.TensorLike, list]: r"""Decomposes any matrix into a linear combination of Pauli operators. This method converts any matrix to a weighted sum of Pauli words acting on :math:`n` qubits diff --git a/pennylane/pauli/dla/center.py b/pennylane/pauli/dla/center.py index b5f518f08a9..6fa5fc66357 100644 --- a/pennylane/pauli/dla/center.py +++ b/pennylane/pauli/dla/center.py @@ -13,7 +13,7 @@ # limitations under the License. """A function to compute the center of a Lie algebra""" from itertools import combinations -from typing import List, Union +from typing import Union import numpy as np @@ -22,8 +22,8 @@ def center( - g: List[Union[Operator, PauliWord, PauliSentence]], pauli: bool = False -) -> List[Union[Operator, PauliSentence]]: + g: list[Union[Operator, PauliWord, PauliSentence]], pauli: bool = False +) -> list[Union[Operator, PauliSentence]]: r""" A function to compute the center of a Lie algebra. diff --git a/pennylane/pauli/dla/lie_closure.py b/pennylane/pauli/dla/lie_closure.py index aadfad83210..931befcfd4f 100644 --- a/pennylane/pauli/dla/lie_closure.py +++ b/pennylane/pauli/dla/lie_closure.py @@ -14,9 +14,10 @@ """A function to compute the Lie closure of a set of operators""" # pylint: disable=too-many-arguments import itertools +from collections.abc import Iterable from copy import copy from functools import reduce -from typing import Iterable, Union +from typing import Union import numpy as np diff --git a/pennylane/pauli/dla/structure_constants.py b/pennylane/pauli/dla/structure_constants.py index f33c58e3c20..e4e9638fc5d 100644 --- a/pennylane/pauli/dla/structure_constants.py +++ b/pennylane/pauli/dla/structure_constants.py @@ -13,7 +13,7 @@ # limitations under the License. """A function to compute the adjoint representation of a Lie algebra""" from itertools import combinations -from typing import List, Union +from typing import Union import numpy as np @@ -34,7 +34,7 @@ def _all_commutators(ops): def structure_constants( - g: List[Union[Operator, PauliWord, PauliSentence]], pauli: bool = False + g: list[Union[Operator, PauliWord, PauliSentence]], pauli: bool = False ) -> TensorLike: r""" Compute the structure constants that make up the adjoint representation of a Lie algebra. diff --git a/pennylane/pauli/utils.py b/pennylane/pauli/utils.py index e6dd201e5c5..7c531572765 100644 --- a/pennylane/pauli/utils.py +++ b/pennylane/pauli/utils.py @@ -22,7 +22,7 @@ """ from functools import lru_cache, singledispatch from itertools import product -from typing import List, Union +from typing import Union import numpy as np @@ -930,7 +930,7 @@ def pauli_group(n_qubits, wire_map=None): @lru_cache() -def partition_pauli_group(n_qubits: int) -> List[List[str]]: +def partition_pauli_group(n_qubits: int) -> list[list[str]]: """Partitions the :math:`n`-qubit Pauli group into qubit-wise commuting terms. The :math:`n`-qubit Pauli group is composed of :math:`4^{n}` terms that can be partitioned into diff --git a/pennylane/pulse/__init__.py b/pennylane/pulse/__init__.py index a5396032ccd..8bbafb5af6d 100644 --- a/pennylane/pulse/__init__.py +++ b/pennylane/pulse/__init__.py @@ -236,7 +236,9 @@ import jax - dev = qml.device("default.qubit.jax", wires=1) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=1) @jax.jit @qml.qnode(dev, interface="jax") @@ -246,10 +248,10 @@ def circuit(params): >>> params = [1.2] >>> circuit(params) -Array(0.96632576, dtype=float32) +Array(0.96632722, dtype=float64) >>> jax.grad(circuit)(params) -[Array(2.3569832, dtype=float32)] +[Array(2.35694829, dtype=float64)] We can use the decorator ``jax.jit`` to compile this execution just-in-time. This means the first execution will typically take a little longer with the benefit that all following executions will be significantly faster. diff --git a/pennylane/pulse/convenience_functions.py b/pennylane/pulse/convenience_functions.py index c1ef567a1c8..dd4f71fabed 100644 --- a/pennylane/pulse/convenience_functions.py +++ b/pennylane/pulse/convenience_functions.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """This file contains convenience functions for pulse programming.""" -from typing import Callable, List, Tuple, Union +from collections.abc import Callable +from typing import Optional, Union import numpy as np @@ -25,7 +26,7 @@ # pylint: disable=unused-argument def constant(scalar, time): - """Returns the given ``scalar``, for use in defining a :class:`~.ParametrizedHamiltonian` with a + r"""Returns the given ``scalar``, for use in defining a :class:`~.ParametrizedHamiltonian` with a trainable coefficient. Args: @@ -57,7 +58,12 @@ def constant(scalar, time): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=1) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit") + @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H)(params, t=2) @@ -66,15 +72,17 @@ def circuit(params): >>> params = jnp.array([5.0]) >>> circuit(params) - Array(0.40808904, dtype=float32) + Array(0.40808193, dtype=float64) >>> jax.grad(circuit)(params) - Array([-3.6517754], dtype=float32) + Array([-3.65178003], dtype=float64) """ return scalar -def rect(x: Union[float, Callable], windows: Union[Tuple[float], List[Tuple[float]]] = None): +def rect( + x: Union[float, Callable], windows: Optional[Union[tuple[float], list[tuple[float]]]] = None +): """Takes a scalar or a scalar-valued function, x, and applies a rectangular window to it, such that the returned function is x inside the window and 0 outside it. diff --git a/pennylane/pulse/hardware_hamiltonian.py b/pennylane/pulse/hardware_hamiltonian.py index e54073f93de..8139d884a8b 100644 --- a/pennylane/pulse/hardware_hamiltonian.py +++ b/pennylane/pulse/hardware_hamiltonian.py @@ -14,8 +14,9 @@ """This module contains the classes/functions needed to simulate and execute the evolution of real Hardware Hamiltonians""" +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable, List, Union +from typing import Optional, Union import numpy as np @@ -95,7 +96,9 @@ def drive(amplitude, phase, wires): .. code-block:: python3 - dev = qml.device("default.qubit.jax", wires=wires) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): @@ -106,7 +109,7 @@ def circuit(params): >>> circuit(params) Array(0.32495208, dtype=float64) >>> jax.grad(circuit)(params) - [Array(1.31956098, dtype=float64)] + [Array(1.31956098, dtype=float64, weak_type=True)] We can also create a Hamiltonian with multiple local drives. The following circuit corresponds to the evolution where an additional local drive that changes in time is acting on wires ``[0, 1]`` is added to the Hamiltonian: @@ -131,12 +134,12 @@ def circuit_local(params): params = (p_global, p_amp, p_phase) >>> circuit_local(params) - Array(-0.5334795, dtype=float64) + Array(0.37385014, dtype=float64) >>> jax.grad(circuit_local)(params) - (Array(0.01654573, dtype=float64), - [Array(-0.04422795, dtype=float64, weak_type=True), - Array(-0.51375441, dtype=float64, weak_type=True)], - Array(0.21901967, dtype=float64)) + (Array(-3.35835837, dtype=float64), + [Array(-3.35835837, dtype=float64, weak_type=True), + Array(-3.35835837, dtype=float64, weak_type=True)], + Array(0.1339487, dtype=float64)) .. details:: :title: Theoretical background @@ -188,7 +191,7 @@ def circuit_local(params): H_d = qml.pulse.drive(amplitude, phase, wires) # detuning term - H_z = qml.dot([-3*np.pi/4]*len(wires), [qml.Z(i) for i in wires]) + H_z = qml.dot([-3*jnp.pi/4]*len(wires), [qml.Z(i) for i in wires]) The total Hamiltonian of that evolution is given by @@ -201,7 +204,7 @@ def circuit_local(params): .. code-block:: python3 - dev = qml.device("default.qubit.jax", wires=wires) + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H_i + H_z + H_d)(params, t=[0, 10]) @@ -209,9 +212,9 @@ def circuit(params): >>> params = [2.4] >>> circuit(params) - Array(0.6962041, dtype=float64) + Array(0.96347734, dtype=float64) >>> jax.grad(circuit)(params) - [Array(1.75825695, dtype=float64)] + [Array(-0.4311521, dtype=float64, weak_type=True)] """ wires = Wires(wires) @@ -306,8 +309,8 @@ def __init__( coeffs, observables, reorder_fn: Callable = _reorder_parameters, - pulses: List["HardwarePulse"] = None, - settings: Union["RydbergSettings", "TransmonSettings"] = None, + pulses: Optional[list["HardwarePulse"]] = None, + settings: Optional[Union["qml.pulse.RydbergSettings", "qml.pulse.TransmonSettings"]] = None, ): self.settings = settings self.pulses = [] if pulses is None else pulses @@ -435,7 +438,7 @@ class HardwarePulse: amplitude: Union[float, Callable] phase: Union[float, Callable] frequency: Union[float, Callable] - wires: List[Wires] + wires: list[Wires] def __post_init__(self): self.wires = Wires(self.wires) diff --git a/pennylane/pulse/parametrized_evolution.py b/pennylane/pulse/parametrized_evolution.py index 3a0dfa08649..fc335b65a80 100644 --- a/pennylane/pulse/parametrized_evolution.py +++ b/pennylane/pulse/parametrized_evolution.py @@ -19,7 +19,8 @@ """ import warnings -from typing import List, Sequence, Union +from collections.abc import Sequence +from typing import Union import pennylane as qml from pennylane.operation import AnyWires, Operation @@ -137,7 +138,9 @@ class ParametrizedEvolution(Operation): import jax - dev = qml.device("default.qubit.jax", wires=1) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=1) @jax.jit @qml.qnode(dev, interface="jax") def circuit(params): @@ -146,10 +149,10 @@ def circuit(params): >>> params = [1.2] >>> circuit(params) - Array(0.96632576, dtype=float32) + Array(0.96632722, dtype=float64) >>> jax.grad(circuit)(params) - [Array(2.3569832, dtype=float32)] + [Array(2.35694829, dtype=float64)] .. note:: In the example above, the decorator ``@jax.jit`` is used to compile this execution just-in-time. This means @@ -223,7 +226,7 @@ def f2(p, t): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=3) + dev = qml.device("default.qubit", wires=3) @qml.qnode(dev, interface="jax") def circuit1(params): @@ -245,11 +248,11 @@ def circuit2(params): >>> params = jnp.array([1., 2., 3.]) >>> circuit1(params) - Array(-0.01543971, dtype=float32) + Array(-0.01542578, dtype=float64) >>> params = jnp.concatenate([params, params]) # H1 + H2 requires 6 parameters! >>> circuit2(params) - Array(-0.78236955, dtype=float32) + Array(-0.78235162, dtype=float64) Here, ``circuit1`` is not executing the evolution of ``H1`` and ``H2`` simultaneously, but rather executing ``H1`` in the ``[0, 10]`` time window and then executing ``H2`` with the same time window, @@ -268,10 +271,10 @@ def circuit(params): return qml.expval(qml.Z(0) @ qml.Z(1) @ qml.Z(2)) >>> circuit(params) - Array(-0.78236955, dtype=float32) + Array(-0.78235162, dtype=float64) >>> jax.grad(circuit)(params) - Array([-4.8066125 , 3.703827 , -1.3297377 , -2.406232 , 0.6811726 , - -0.52277344], dtype=float32) + Array([-4.80708632, 3.70323783, -1.32958799, -2.40642477, 0.68105214, + -0.52269657], dtype=float64) Given that we used the same time window (``[0, 10]``), the results are the same as before. @@ -311,7 +314,7 @@ def circuit(params): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=1) + dev = qml.device("default.qubit", wires=1) @qml.qnode(dev, interface="jax") def circuit(param, time): @@ -320,11 +323,11 @@ def circuit(param, time): >>> circuit(param, time) Array([[1. , 0. ], - [0.9897738 , 0.01022595], - [0.9599043 , 0.04009585], - [0.9123617 , 0.08763832], - [0.84996957, 0.15003097], - [0.7761489 , 0.22385144]], dtype=float32) + [0.98977406, 0.01022594], + [0.95990416, 0.04009584], + [0.91236167, 0.08763833], + [0.84996865, 0.15003133], + [0.77614817, 0.22385181]], dtype=float64) **Computing complementary time evolution** @@ -371,7 +374,7 @@ def __init__( self, H: ParametrizedHamiltonian, params: list = None, - t: Union[float, List[float]] = None, + t: Union[float, list[float]] = None, return_intermediate: bool = False, complementary: bool = False, dense: bool = None, diff --git a/pennylane/pulse/parametrized_hamiltonian_pytree.py b/pennylane/pulse/parametrized_hamiltonian_pytree.py index 9e05f174519..958dc2c11ce 100644 --- a/pennylane/pulse/parametrized_hamiltonian_pytree.py +++ b/pennylane/pulse/parametrized_hamiltonian_pytree.py @@ -12,8 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. """Module containing the ``JaxParametrizedHamiltonian`` class.""" +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable, Optional, Tuple, Union +from typing import Optional, Union import jax import jax.numpy as jnp @@ -32,8 +33,8 @@ class ParametrizedHamiltonianPytree: """Jax pytree class that represents a ``ParametrizedHamiltonian``.""" mat_fixed: Optional[Union[jnp.ndarray, sparse.BCSR]] - mats_parametrized: Tuple[Union[jnp.ndarray, sparse.BCSR], ...] - coeffs_parametrized: Tuple[Callable] + mats_parametrized: tuple[Union[jnp.ndarray, sparse.BCSR], ...] + coeffs_parametrized: tuple[Callable] reorder_fn: Callable @staticmethod @@ -117,8 +118,8 @@ def tree_unflatten(cls, param_coeffs: tuple, matrices: tuple, reorder_fn: callab class LazyDotPytree: """Jax pytree representing a lazy dot operation.""" - coeffs: Tuple[complex, ...] - mats: Tuple[Union[jnp.ndarray, sparse.BCSR], ...] + coeffs: tuple[complex, ...] + mats: tuple[Union[jnp.ndarray, sparse.BCSR], ...] @jax.jit def __matmul__(self, other): diff --git a/pennylane/pulse/rydberg.py b/pennylane/pulse/rydberg.py index 761c61ffdc8..2059d3e9fd5 100644 --- a/pennylane/pulse/rydberg.py +++ b/pennylane/pulse/rydberg.py @@ -86,7 +86,11 @@ def rydberg_interaction( .. code-block:: python - dev = qml.device("default.qubit.jax", wires=9) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=9) @qml.qnode(dev, interface="jax") def circuit(): @@ -94,7 +98,7 @@ def circuit(): return qml.expval(qml.Z(0)) >>> circuit() - Array(1., dtype=float32) + Array(1., dtype=float64) """ if wires is None: wires = list(range(len(register))) @@ -208,7 +212,11 @@ def rydberg_drive(amplitude, phase, detuning, wires): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=wires) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H_i + H_d)(params, t=[0, 0.5]) diff --git a/pennylane/pulse/transmon.py b/pennylane/pulse/transmon.py index 15662340764..7262d54f6ce 100644 --- a/pennylane/pulse/transmon.py +++ b/pennylane/pulse/transmon.py @@ -13,8 +13,9 @@ # limitations under the License. """This module contains the classes/functions specific for simulation of superconducting transmon hardware systems""" import warnings +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable, List, Union +from typing import Union import numpy as np @@ -225,7 +226,7 @@ class TransmonSettings: """ - connections: List + connections: list qubit_freq: Union[float, Callable] coupling: Union[list, TensorLike, Callable] anharmonicity: Union[float, Callable] @@ -350,6 +351,10 @@ def phase(phi0, t): .. code-block:: python3 + import jax + + jax.config.update("jax_enable_x64", True) + qubit_freqs = [5.1, 5., 5.3] connections = [[0, 1], [1, 2]] # qubits 0 and 1 are coupled, as are 1 and 2 g = [0.02, 0.05] @@ -363,7 +368,7 @@ def amp(max_amp, t): return max_amp * jnp.sin(t) ** 2 for q in range(3): H += qml.pulse.transmon_drive(amp, phase, freq, q) # Parametrized drive for each qubit - dev = qml.device("default.qubit.jax", wires=range(3)) + dev = qml.device("default.qubit", wires=range(3)) @jax.jit @qml.qnode(dev, interface="jax") diff --git a/pennylane/qaoa/cost.py b/pennylane/qaoa/cost.py index 7a67af6baa7..7ffc65be5e7 100644 --- a/pennylane/qaoa/cost.py +++ b/pennylane/qaoa/cost.py @@ -16,7 +16,8 @@ different optimization problems. """ # pylint: disable=unnecessary-lambda-assignment -from typing import Iterable, Union +from collections.abc import Iterable +from typing import Union import networkx as nx import rustworkx as rx diff --git a/pennylane/qaoa/cycle.py b/pennylane/qaoa/cycle.py index 1fef9d74977..0edfb4b9146 100644 --- a/pennylane/qaoa/cycle.py +++ b/pennylane/qaoa/cycle.py @@ -16,7 +16,8 @@ """ # pylint: disable=unnecessary-comprehension, unnecessary-lambda-assignment import itertools -from typing import Dict, Iterable, List, Tuple, Union +from collections.abc import Iterable +from typing import Union import networkx as nx import numpy as np @@ -25,7 +26,7 @@ import pennylane as qml -def edges_to_wires(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> Dict[Tuple, int]: +def edges_to_wires(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> dict[tuple, int]: r"""Maps the edges of a graph to corresponding wires. **Example** @@ -79,7 +80,7 @@ def edges_to_wires(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> Dict[Tup ) -def wires_to_edges(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> Dict[int, Tuple]: +def wires_to_edges(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> dict[int, tuple]: r"""Maps the wires of a register of qubits to corresponding edges. **Example** @@ -230,7 +231,7 @@ def cycle_mixer(graph: Union[nx.DiGraph, rx.PyDiGraph]) -> qml.operation.Operato def _partial_cycle_mixer( - graph: Union[nx.DiGraph, rx.PyDiGraph], edge: Tuple + graph: Union[nx.DiGraph, rx.PyDiGraph], edge: tuple ) -> qml.operation.Operator: r"""Calculates the partial cycle-mixer Hamiltonian for a specific edge. @@ -414,7 +415,7 @@ def loss_hamiltonian(graph: Union[nx.Graph, rx.PyGraph, rx.PyDiGraph]) -> qml.op def _square_hamiltonian_terms( coeffs: Iterable[float], ops: Iterable[qml.operation.Observable] -) -> Tuple[List[float], List[qml.operation.Observable]]: +) -> tuple[list[float], list[qml.operation.Observable]]: """Calculates the coefficients and observables that compose the squared Hamiltonian. Args: diff --git a/pennylane/qaoa/layers.py b/pennylane/qaoa/layers.py index 0c1f7ab9c1d..dde6a39a998 100644 --- a/pennylane/qaoa/layers.py +++ b/pennylane/qaoa/layers.py @@ -91,7 +91,7 @@ def circuit(gamma): >>> print(qml.draw(circuit)(0.5)) 0: ──H─╭ApproxTimeEvolution(1.00,1.00,0.50)── 1: ──H─╰ApproxTimeEvolution(1.00,1.00,0.50)── - >>> print(qml.draw(circuit, expansion_strategy="device")(0.5)) + >>> print(qml.draw(circuit, level="device")(0.5)) 0: ──H──RZ(1.00)─╭RZZ(1.00)── 1: ──H───────────╰RZZ(1.00)── @@ -149,7 +149,7 @@ def circuit(alpha): >>> print(qml.draw(circuit)(0.5)) 0: ──H─╭ApproxTimeEvolution(1.00,1.00,0.50)── 1: ──H─╰ApproxTimeEvolution(1.00,1.00,0.50)── - >>> print(qml.draw(circuit, expansion_strategy="device")(0.5)) + >>> print(qml.draw(circuit, level="device")(0.5)) 0: ──H──RX(1.00)─╭RXX(1.00)── 1: ──H───────────╰RXX(1.00)── diff --git a/pennylane/qaoa/mixers.py b/pennylane/qaoa/mixers.py index 7c2a70018f2..0e96c0da991 100644 --- a/pennylane/qaoa/mixers.py +++ b/pennylane/qaoa/mixers.py @@ -18,7 +18,8 @@ # pylint: disable=unnecessary-lambda-assignment import itertools -from typing import Iterable, Union +from collections.abc import Iterable +from typing import Union import networkx as nx import rustworkx as rx diff --git a/pennylane/qchem/observable_hf.py b/pennylane/qchem/observable_hf.py index 91b100de54f..276bfac9104 100644 --- a/pennylane/qchem/observable_hf.py +++ b/pennylane/qchem/observable_hf.py @@ -20,6 +20,7 @@ import pennylane as qml from pennylane.fermi import FermiSentence, FermiWord from pennylane.operation import active_new_opmath +from pennylane.pauli import PauliSentence from pennylane.pauli.utils import simplify @@ -81,12 +82,9 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12): coeffs = qml.math.concatenate((coeffs, coeffs_two)) operators = operators + operators_two - indices_sort = [operators.index(i) for i in sorted(operators)] - if indices_sort: - indices_sort = qml.math.array(indices_sort) - sentence = FermiSentence({FermiWord({}): constant[0]}) - for c, o in zip(coeffs[indices_sort], sorted(operators)): + for c, o in sorted(zip(coeffs, operators), key=lambda item: item[1]): + if len(o) == 2: sentence.update({FermiWord({(0, o[0]): "+", (1, o[1]): "-"}): c}) if len(o) == 4: @@ -145,6 +143,11 @@ def qubit_observable(o_ferm, cutoff=1.0e-12, mapping="jordan_wigner"): qubits = len(o_ferm.wires) h = qml.bravyi_kitaev(o_ferm, qubits, ps=True, tol=cutoff) + if list(h.wires) != sorted(list(h.wires)): + h = PauliSentence( + sorted(h.items(), key=lambda item: max(item[0].wires.tolist(), default=0)) + ) + h.simplify(tol=cutoff) if active_new_opmath(): diff --git a/pennylane/qcut/cutcircuit.py b/pennylane/qcut/cutcircuit.py index 2c9a1c862ac..60740cc96d9 100644 --- a/pennylane/qcut/cutcircuit.py +++ b/pennylane/qcut/cutcircuit.py @@ -15,13 +15,15 @@ Function cut_circuit for cutting a quantum circuit into smaller circuit fragments. """ +from collections.abc import Callable from functools import partial -from typing import Callable, Optional, Sequence, Union +from typing import Optional, Union import pennylane as qml from pennylane.measurements import ExpectationMP -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires from .cutstrategy import CutStrategy @@ -38,7 +40,7 @@ def _cut_circuit_expand( max_depth: int = 1, auto_cutter: Union[bool, Callable] = False, **kwargs, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Main entry point for expanding operations until reaching a depth that includes :class:`~.WireCut` operations.""" # pylint: disable=unused-argument @@ -63,7 +65,7 @@ def processing_fn(res): tape.operations, [new_meas_op], shots=tape.shots, trainable_params=tape.trainable_params ) - tapes, tapes_fn = qml.transforms.hamiltonian_expand(new_tape, group=False) + tapes, tapes_fn = qml.transforms.split_non_commuting(new_tape, grouping_strategy=None) return [_qcut_expand_fn(tape, max_depth, auto_cutter) for tape in tapes], tapes_fn @@ -76,7 +78,7 @@ def cut_circuit( device_wires: Optional[Wires] = None, max_depth: int = 1, **kwargs, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """ Cut up a quantum circuit into smaller circuit fragments. diff --git a/pennylane/qcut/cutcircuit_mc.py b/pennylane/qcut/cutcircuit_mc.py index d7c5f48c72e..57580645cef 100644 --- a/pennylane/qcut/cutcircuit_mc.py +++ b/pennylane/qcut/cutcircuit_mc.py @@ -16,16 +16,18 @@ Monte Carlo method, at its auxillary functions""" import inspect +from collections.abc import Callable from functools import partial -from typing import Callable, List, Optional, Sequence, Tuple, Union +from typing import Optional, Union import numpy as np from networkx import MultiDiGraph import pennylane as qml from pennylane.measurements import SampleMP -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires from .cutstrategy import CutStrategy @@ -55,7 +57,7 @@ def _cut_circuit_mc_expand( device_wires: Optional[Wires] = None, auto_cutter: Union[bool, Callable] = False, **kwargs, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Main entry point for expanding operations in sample-based tapes until reaching a depth that includes :class:`~.WireCut` operations.""" # pylint: disable=unused-argument, too-many-arguments @@ -75,7 +77,7 @@ def cut_circuit_mc( shots: Optional[int] = None, device_wires: Optional[Wires] = None, **kwargs, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """ Cut up a circuit containing sample measurements into smaller fragments using a Monte Carlo method. @@ -591,8 +593,8 @@ def _pauliZ(wire): def expand_fragment_tapes_mc( - tapes: Sequence[QuantumTape], communication_graph: MultiDiGraph, shots: int -) -> Tuple[List[QuantumTape], np.ndarray]: + tapes: QuantumTapeBatch, communication_graph: MultiDiGraph, shots: int +) -> tuple[QuantumTapeBatch, np.ndarray]: """ Expands fragment tapes into a sequence of random configurations of the contained pairs of :class:`MeasureNode` and :class:`PrepareNode` operations. @@ -617,7 +619,7 @@ def expand_fragment_tapes_mc( shots (int): number of shots Returns: - Tuple[List[QuantumTape], np.ndarray]: the tapes corresponding to each configuration and the + Tuple[Sequence[QuantumTape], np.ndarray]: the tapes corresponding to each configuration and the settings that track each configuration pair **Example** diff --git a/pennylane/qcut/cutstrategy.py b/pennylane/qcut/cutstrategy.py index 7406b6bafab..a0998e8a8bb 100644 --- a/pennylane/qcut/cutstrategy.py +++ b/pennylane/qcut/cutstrategy.py @@ -16,9 +16,9 @@ """ import warnings -from collections.abc import Sequence as SequenceType +from collections.abc import Sequence from dataclasses import InitVar, dataclass -from typing import Any, ClassVar, Dict, List, Sequence, Union +from typing import Any, ClassVar, Union from networkx import MultiDiGraph @@ -129,7 +129,7 @@ def __post_init__( devices = (devices,) if devices is not None: - if not isinstance(devices, SequenceType) or any( + if not isinstance(devices, Sequence) or any( (not isinstance(d, (qml.devices.LegacyDevice, qml.devices.Device)) for d in devices) ): raise ValueError( @@ -158,7 +158,7 @@ def get_cut_kwargs( max_wires_by_fragment: Sequence[int] = None, max_gates_by_fragment: Sequence[int] = None, exhaustive: bool = True, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Derive the complete set of arguments, based on a given circuit, for passing to a graph partitioner. @@ -284,7 +284,7 @@ def _infer_probed_cuts( max_wires_by_fragment=None, max_gates_by_fragment=None, exhaustive=True, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """ Helper function for deriving the minimal set of best default partitioning constraints for the graph partitioner. diff --git a/pennylane/qcut/kahypar.py b/pennylane/qcut/kahypar.py index 8ba7a0743a4..20603b8e806 100644 --- a/pennylane/qcut/kahypar.py +++ b/pennylane/qcut/kahypar.py @@ -16,10 +16,10 @@ """ -from collections.abc import Sequence as SequenceType +from collections.abc import Sequence from itertools import compress from pathlib import Path -from typing import Any, List, Sequence, Tuple, Union +from typing import Any, Union import numpy as np from networkx import MultiDiGraph @@ -32,15 +32,15 @@ def kahypar_cut( graph: MultiDiGraph, num_fragments: int, imbalance: int = None, - edge_weights: List[Union[int, float]] = None, - node_weights: List[Union[int, float]] = None, - fragment_weights: List[Union[int, float]] = None, + edge_weights: list[Union[int, float]] = None, + node_weights: list[Union[int, float]] = None, + fragment_weights: list[Union[int, float]] = None, hyperwire_weight: int = 1, seed: int = None, config_path: Union[str, Path] = None, trial: int = None, verbose: bool = False, -) -> List[Tuple[Operation, Operation, Any]]: +) -> list[tuple[Operation, Any]]: """Calls `KaHyPar `__ to partition a graph. .. warning:: @@ -141,7 +141,7 @@ def kahypar_cut( if isinstance(imbalance, float): context.setEpsilon(imbalance) - if isinstance(fragment_weights, SequenceType) and (len(fragment_weights) == num_fragments): + if isinstance(fragment_weights, Sequence) and (len(fragment_weights) == num_fragments): context.setCustomTargetBlockWeights(fragment_weights) if not verbose: context.suppressOutput(True) @@ -168,7 +168,7 @@ def _graph_to_hmetis( graph: MultiDiGraph, hyperwire_weight: int = 0, edge_weights: Sequence[int] = None, -) -> Tuple[List[int], List[int], List[Union[int, float]]]: +) -> tuple[list[int], list[int], list[Union[int, float]]]: """Converts a ``MultiDiGraph`` into the `hMETIS hypergraph input format `__ conforming to KaHyPar's calling signature. diff --git a/pennylane/qcut/processing.py b/pennylane/qcut/processing.py index 0941cbbc341..4593bf2a3cf 100644 --- a/pennylane/qcut/processing.py +++ b/pennylane/qcut/processing.py @@ -16,7 +16,7 @@ """ import string -from typing import List, Sequence +from collections.abc import Sequence from networkx import MultiDiGraph @@ -84,7 +84,7 @@ def qcut_processing_fn( def qcut_processing_fn_sample( results: Sequence, communication_graph: MultiDiGraph, shots: int -) -> List: +) -> list: """ Function to postprocess samples for the :func:`cut_circuit_mc() ` transform. This removes superfluous mid-circuit measurement samples from fragment @@ -187,7 +187,7 @@ def qcut_processing_fn_mc( return qml.math.convert_like(pnp.mean(expvals), res0) -def _reshape_results(results: Sequence, shots: int) -> List[List]: +def _reshape_results(results: Sequence, shots: int) -> list[list]: """ Helper function to reshape ``results`` into a two-dimensional nested list whose number of rows is determined by the number of shots and whose number of columns is determined by the number of @@ -428,7 +428,7 @@ def _to_tensors( results, prepare_nodes: Sequence[Sequence[PrepareNode]], measure_nodes: Sequence[Sequence[MeasureNode]], -) -> List: +) -> list: """Process a flat list of execution results from all circuit fragments into the corresponding tensors. diff --git a/pennylane/qcut/tapes.py b/pennylane/qcut/tapes.py index 9765d088430..0d9755120c4 100644 --- a/pennylane/qcut/tapes.py +++ b/pennylane/qcut/tapes.py @@ -16,8 +16,9 @@ """ import copy +from collections.abc import Callable, Sequence from itertools import product -from typing import Callable, List, Sequence, Tuple, Union +from typing import Union from networkx import MultiDiGraph @@ -258,7 +259,7 @@ def _prep_iplus(wire): def expand_fragment_tape( tape: QuantumTape, -) -> Tuple[List[QuantumTape], List[PrepareNode], List[MeasureNode]]: +) -> tuple[list[QuantumTape], list[PrepareNode], list[MeasureNode]]: """ Expands a fragment tape into a sequence of tapes for each configuration of the contained :class:`MeasureNode` and :class:`PrepareNode` operations. @@ -354,7 +355,7 @@ def expand_fragment_tape( def _get_measurements( group: Sequence[Operator], measurements: Sequence[MeasurementProcess] -) -> List[MeasurementProcess]: +) -> list[MeasurementProcess]: """Pairs each observable in ``group`` with the circuit ``measurements``. Only a single measurement of an expectation value is currently supported diff --git a/pennylane/qcut/utils.py b/pennylane/qcut/utils.py index 4aafa8c236c..11068e5a86b 100644 --- a/pennylane/qcut/utils.py +++ b/pennylane/qcut/utils.py @@ -18,7 +18,8 @@ import uuid import warnings -from typing import Any, Callable, Sequence, Tuple +from collections.abc import Callable, Sequence +from typing import Any import numpy as np from networkx import MultiDiGraph, has_path, weakly_connected_components @@ -466,7 +467,7 @@ def replace_wire_cut_nodes(graph: MultiDiGraph): def place_wire_cuts( - graph: MultiDiGraph, cut_edges: Sequence[Tuple[Operation, Operation, Any]] + graph: MultiDiGraph, cut_edges: Sequence[tuple[Operation, Operation, Any]] ) -> MultiDiGraph: """Inserts a :class:`~.WireCut` node for each provided cut edge into a circuit graph. @@ -571,7 +572,7 @@ def _remove_existing_cuts(graph: MultiDiGraph) -> MultiDiGraph: # pylint: disable=too-many-branches -def fragment_graph(graph: MultiDiGraph) -> Tuple[Tuple[MultiDiGraph], MultiDiGraph]: +def fragment_graph(graph: MultiDiGraph) -> tuple[tuple[MultiDiGraph], MultiDiGraph]: """ Fragments a graph into a collection of subgraphs as well as returning the communication (`quotient `__) diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py index 6edaf8c3c04..1aa51cb8afa 100644 --- a/pennylane/qinfo/transforms.py +++ b/pennylane/qinfo/transforms.py @@ -14,19 +14,20 @@ """QNode transforms for the quantum information quantities.""" # pylint: disable=import-outside-toplevel, not-callable import warnings +from collections.abc import Callable, Sequence from functools import partial -from typing import Callable, Sequence import pennylane as qml from pennylane import transform from pennylane.devices import DefaultMixed, DefaultQubit, DefaultQubitLegacy from pennylane.gradients import adjoint_metric_tensor, metric_tensor from pennylane.measurements import DensityMatrixMP, StateMP -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn @partial(transform, final_transform=True) -def reduced_dm(tape: QuantumTape, wires, **kwargs) -> (Sequence[QuantumTape], Callable): +def reduced_dm(tape: QuantumTape, wires, **kwargs) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Compute the reduced density matrix from a :class:`~.QNode` returning :func:`~pennylane.state`. @@ -123,7 +124,7 @@ def _reduced_dm_qnode(self, qnode, targs, tkwargs): @partial(transform, final_transform=True) -def purity(tape: QuantumTape, wires, **kwargs) -> (Sequence[QuantumTape], Callable): +def purity(tape: QuantumTape, wires, **kwargs) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Compute the purity of a :class:`~.QuantumTape` returning :func:`~pennylane.state`. .. math:: @@ -222,7 +223,7 @@ def _purity_qnode(self, qnode, targs, tkwargs): @partial(transform, final_transform=True) def vn_entropy( tape: QuantumTape, wires: Sequence[int], base: float = None, **kwargs -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Compute the Von Neumann entropy from a :class:`.QuantumTape` returning a :func:`~pennylane.state`. .. math:: @@ -352,7 +353,7 @@ def processing_fn(res): @partial(transform, final_transform=True) def mutual_info( tape: QuantumTape, wires0: Sequence[int], wires1: Sequence[int], base: float = None, **kwargs -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Compute the mutual information from a :class:`.QuantumTape` returning a :func:`~pennylane.state`: .. math:: @@ -617,7 +618,7 @@ def circ(params): @partial(transform, is_informative=True) def quantum_fisher( tape: qml.tape.QuantumTape, device, *args, **kwargs -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Returns a function that computes the quantum fisher information matrix (QFIM) of a given :class:`.QNode`. Given a parametrized quantum state :math:`|\psi(\bm{\theta})\rangle`, the quantum fisher information matrix (QFIM) quantifies how changes to the parameters :math:`\bm{\theta}` diff --git a/pennylane/qnn/torch.py b/pennylane/qnn/torch.py index bbeef534a5e..0fb0efc39fe 100644 --- a/pennylane/qnn/torch.py +++ b/pennylane/qnn/torch.py @@ -18,8 +18,8 @@ import functools import inspect import math -from collections.abc import Iterable -from typing import Any, Callable, Dict, Union +from collections.abc import Callable, Iterable +from typing import Any, Union from pennylane import QNode @@ -328,7 +328,7 @@ def __init__( self, qnode: QNode, weight_shapes: dict, - init_method: Union[Callable, Dict[str, Union[Callable, Any]]] = None, + init_method: Union[Callable, dict[str, Union[Callable, Any]]] = None, # FIXME: Cannot change type `Any` to `torch.Tensor` in init_method because it crashes the # tests that don't use torch module. ): @@ -354,7 +354,7 @@ def __init__( if self.qnode.interface not in ("auto", "torch", "pytorch"): raise ValueError(f"Invalid interface '{self.qnode.interface}' for TorchLayer") - self.qnode_weights: Dict[str, torch.nn.Parameter] = {} + self.qnode_weights: dict[str, torch.nn.Parameter] = {} self._init_weights(init_method=init_method, weight_shapes=weight_shapes) self._initialized = True @@ -480,8 +480,8 @@ def __setattr__(self, item, val): def _init_weights( self, - weight_shapes: Dict[str, tuple], - init_method: Union[Callable, Dict[str, Union[Callable, Any]], None], + weight_shapes: dict[str, tuple], + init_method: Union[Callable, dict[str, Union[Callable, Any]], None], ): r"""Initialize and register the weights with the given init_method. If init_method is not specified, weights are randomly initialized from the uniform distribution on the interval diff --git a/pennylane/resource/error/error.py b/pennylane/resource/error/error.py index 53a5ca39c9d..59d885e76b9 100644 --- a/pennylane/resource/error/error.py +++ b/pennylane/resource/error/error.py @@ -15,7 +15,6 @@ Stores classes and logic to define and track algorithmic error in a quantum workflow. """ from abc import ABC, abstractmethod -from typing import Dict import pennylane as qml from pennylane.operation import Operation, Operator @@ -145,7 +144,7 @@ def get_error(approximate_op: Operator, exact_op: Operator): return qml.math.max(qml.math.svd(m1 - m2, compute_uv=False)) -def _compute_algo_error(tape) -> Dict[str, AlgorithmicError]: +def _compute_algo_error(tape) -> dict[str, AlgorithmicError]: """Given a quantum circuit (tape), this function computes the algorithmic error generated by standard PennyLane operations. diff --git a/pennylane/resource/specs.py b/pennylane/resource/specs.py index 371af2fff09..95cc3221a64 100644 --- a/pennylane/resource/specs.py +++ b/pennylane/resource/specs.py @@ -25,7 +25,8 @@ def _get_absolute_import_path(fn): def _determine_spec_level(kwargs, qnode): if "max_expansion" in kwargs: warnings.warn( - "'max_expansion' has no effect on the output of 'specs()' and should not be used." + "'max_expansion' has no effect on the output of 'specs()' and should not be used.", + qml.PennyLaneDeprecationWarning, ) sentinel = object() @@ -36,6 +37,14 @@ def _determine_spec_level(kwargs, qnode): if all(val != sentinel for val in (level, expansion_strategy)): raise ValueError("Either 'level' or 'expansion_strategy' need to be set, but not both.") + if expansion_strategy != sentinel: + warnings.warn( + "The 'expansion_strategy' argument is deprecated and will be removed in " + "version 0.39. Instead, use the 'level' argument which offers more flexibility " + "and options.", + qml.PennyLaneDeprecationWarning, + ) + if level == sentinel: if expansion_strategy == sentinel: return qnode.expansion_strategy @@ -86,6 +95,9 @@ def specs(qnode, **kwargs): ``max_expansion`` and ``qnode.max_expansion`` have no effect on the return of this function and will be ignored. + .. warning:: + The ``expansion_strategy`` argument is deprecated and will be removed in version 0.39. Use the ``level`` + argument instead to specify the resulting tape you want. **Example** @@ -198,7 +210,7 @@ def circuit(x): H = qml.Hamiltonian([0.2, -0.543], [qml.X(0) @ qml.Z(1), qml.Z(0) @ qml.Y(2)]) - @qml.transforms.hamiltonian_expand + @qml.transforms.split_non_commuting @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=np.pi / 4) def circuit(): qml.RandomLayers(qml.numpy.array([[1.0, 2.0]]), wires=(0, 1)) diff --git a/pennylane/shadows/transforms.py b/pennylane/shadows/transforms.py index 76a136a1606..f0a2e234e0f 100644 --- a/pennylane/shadows/transforms.py +++ b/pennylane/shadows/transforms.py @@ -16,18 +16,20 @@ import warnings from functools import partial, reduce from itertools import product -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane import transform from pennylane.measurements import ClassicalShadowMP -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch +from pennylane.typing import PostprocessingFn @transform -def _replace_obs(tape: QuantumTape, obs, *args, **kwargs) -> (Sequence[QuantumTape], Callable): +def _replace_obs( + tape: QuantumTape, obs, *args, **kwargs +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """ Tape transform to replace the measurement processes with the given one """ @@ -54,7 +56,7 @@ def processing_fn(res): @partial(transform, final_transform=True) -def shadow_expval(tape: QuantumTape, H, k=1) -> (Sequence[QuantumTape], Callable): +def shadow_expval(tape: QuantumTape, H, k=1) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transform a circuit returning a classical shadow into one that returns the approximate expectation values in a differentiable manner. @@ -170,7 +172,9 @@ def post_processing(results): @partial(transform, final_transform=True) -def shadow_state(tape: QuantumTape, wires, diffable=False) -> (Sequence[QuantumTape], Callable): +def shadow_state( + tape: QuantumTape, wires, diffable=False +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transform a circuit returning a classical shadow into one that returns the reconstructed state in a differentiable manner. diff --git a/pennylane/tape/__init__.py b/pennylane/tape/__init__.py index aebe25ba8cd..546853e280d 100644 --- a/pennylane/tape/__init__.py +++ b/pennylane/tape/__init__.py @@ -18,4 +18,4 @@ from .operation_recorder import OperationRecorder from .qscript import QuantumScript, make_qscript -from .tape import QuantumTape, TapeError, expand_tape_state_prep +from .tape import QuantumTape, QuantumTapeBatch, TapeError, expand_tape_state_prep diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py index dbd8e2aa7ff..271929ad8c8 100644 --- a/pennylane/tape/qscript.py +++ b/pennylane/tape/qscript.py @@ -20,8 +20,9 @@ import contextlib import copy from collections import Counter +from collections.abc import Sequence from functools import cached_property -from typing import List, Optional, Sequence, Union +from typing import Optional, Union import pennylane as qml from pennylane.measurements import MeasurementProcess, ProbabilityMP, Shots, StateMP @@ -238,7 +239,7 @@ def circuit(self): return self.operations + self.measurements @property - def operations(self) -> List[Operator]: + def operations(self) -> list[Operator]: """Returns the state preparations and operations on the quantum script. Returns: @@ -252,7 +253,7 @@ def operations(self) -> List[Operator]: return self._ops @property - def observables(self) -> List[Union[MeasurementProcess, Observable]]: + def observables(self) -> list[Union[MeasurementProcess, Observable]]: """Returns the observables on the quantum script. Returns: @@ -279,7 +280,7 @@ def observables(self) -> List[Union[MeasurementProcess, Observable]]: return obs @property - def measurements(self) -> List[MeasurementProcess]: + def measurements(self) -> list[MeasurementProcess]: """Returns the measurements on the quantum script. Returns: @@ -326,7 +327,7 @@ def output_dim(self): return self._output_dim @property - def diagonalizing_gates(self) -> List[Operation]: + def diagonalizing_gates(self) -> list[Operation]: """Returns the gates that diagonalize the measured wires such that they are in the eigenbasis of the circuit observables. @@ -537,7 +538,7 @@ def data(self): @property def trainable_params(self): - """Store or return a list containing the indices of parameters that support + r"""Store or return a list containing the indices of parameters that support differentiability. The indices provided match the order of appearence in the quantum circuit. diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py index c0f1ba68091..92a4bbfbe90 100644 --- a/pennylane/tape/tape.py +++ b/pennylane/tape/tape.py @@ -16,6 +16,7 @@ """ # pylint: disable=too-many-instance-attributes,protected-access,too-many-branches,too-many-public-methods, too-many-arguments import copy +from collections.abc import Sequence from threading import RLock import pennylane as qml @@ -541,4 +542,6 @@ def __hash__(self): return QuantumScript.__hash__(self) +QuantumTapeBatch = Sequence[QuantumTape] + register_pytree(QuantumTape, QuantumTape._flatten, QuantumTape._unflatten) diff --git a/pennylane/templates/embeddings/angle.py b/pennylane/templates/embeddings/angle.py index d01ef8f3170..6cb7eafd6d5 100644 --- a/pennylane/templates/embeddings/angle.py +++ b/pennylane/templates/embeddings/angle.py @@ -65,7 +65,7 @@ def circuit(feature_vector): Here, we have also used rotation angles :class:`RZ`. If not specified, :class:`RX` is used as default. The resulting circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")(X)) + >>> print(qml.draw(circuit, level="device")(X)) 0: ──RZ(1.00)──H── β•­Probs 1: ──RZ(2.00)───── β”œProbs 2: ──RZ(3.00)───── β•°Probs diff --git a/pennylane/templates/embeddings/basis.py b/pennylane/templates/embeddings/basis.py index a812c644ca3..18175a8fd6a 100644 --- a/pennylane/templates/embeddings/basis.py +++ b/pennylane/templates/embeddings/basis.py @@ -55,7 +55,7 @@ def circuit(feature_vector): The resulting circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")(X)) + >>> print(qml.draw(circuit, level="device")(X)) 0: ──X── State 1: ──X── State 2: ──X── State diff --git a/pennylane/templates/layers/random.py b/pennylane/templates/layers/random.py index e00bb0886fd..c36e649ae31 100644 --- a/pennylane/templates/layers/random.py +++ b/pennylane/templates/layers/random.py @@ -91,11 +91,11 @@ def circuit2(weights): You can verify this by drawing the circuits. - >>> print(qml.draw(circuit1, expansion_strategy="device")(weights)) + >>> print(qml.draw(circuit1, level="device")(weights)) 0: ──────────────────────╭X─╭X──RZ(1.40)── 1: ──RX(0.10)──RX(-2.10)─╰●─╰●──────────── - >>> print(qml.draw(circuit2, expansion_strategy="device")(weights)) + >>> print(qml.draw(circuit2, level="device")(weights)) 0: ──────────────────────╭X─╭X──RZ(1.40)── 1: ──RX(0.10)──RX(-2.10)─╰●─╰●──────────── @@ -111,10 +111,10 @@ def circuit2(weights): ... return qml.expval(qml.Z(0)) >>> np.allclose(circuit(weights, seed=9), circuit(weights, seed=12)) False - >>> print(qml.draw(circuit, expansion_strategy="device")(weights, seed=9)) + >>> print(qml.draw(circuit, level="device")(weights, seed=9)) 0: ─╭X──RX(0.10)───────────── 1: ─╰●──RY(-2.10)──RX(1.40)── - >>> print(qml.draw(circuit, expansion_strategy="device")(weights, seed=12)) + >>> print(qml.draw(circuit, level="device")(weights, seed=12)) 0: ─╭X──RZ(0.10)──╭●─╭X──────────── 1: ─╰●──RX(-2.10)─╰X─╰●──RZ(1.40)── diff --git a/pennylane/templates/layers/strongly_entangling.py b/pennylane/templates/layers/strongly_entangling.py index 554f100e8f5..7ffa6ae8a13 100644 --- a/pennylane/templates/layers/strongly_entangling.py +++ b/pennylane/templates/layers/strongly_entangling.py @@ -72,7 +72,7 @@ def circuit(parameters): The resulting circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")(weights)) + >>> print(qml.draw(circuit, level="device")(weights)) 0: ──Rot(0.68,0.98,0.48)─╭●───────╭X──Rot(0.94,0.22,0.70)─╭●────╭X───── 1: ──Rot(0.91,0.19,0.15)─╰X─╭●────│───Rot(0.50,0.20,0.63)─│──╭●─│──╭X── 2: ──Rot(0.91,0.68,0.96)────╰X─╭●─│───Rot(0.14,0.05,0.16)─╰X─│──╰●─│─── @@ -108,7 +108,7 @@ def circuit(parameters): The resulting circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")(weights)) + >>> print(qml.draw(circuit, level="device")(weights)) 0: ──Rot(0.99,0.17,0.12)─╭●────╭Z──Rot(0.02,0.94,0.57)──────────────────────╭●─╭Z──────── 1: ──Rot(0.55,0.42,0.61)─│──╭●─│──╭Z────────────────────Rot(0.15,0.26,0.82)─│──╰●─╭Z───── 2: ──Rot(0.79,0.93,0.27)─╰Z─│──╰●─│─────────────────────Rot(0.73,0.01,0.44)─│─────╰●─╭Z── diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py index aaaece40046..b86092c93af 100644 --- a/pennylane/templates/state_preparations/mottonen.py +++ b/pennylane/templates/state_preparations/mottonen.py @@ -267,7 +267,7 @@ def circuit(state): state = np.array([1, 2j, 3, 4j, 5, 6j, 7, 8j]) state = state / np.linalg.norm(state) - print(qml.draw(circuit, expansion_strategy="device", max_length=80)(state)) + print(qml.draw(circuit, level="device", max_length=80)(state)) .. code-block:: diff --git a/pennylane/templates/subroutines/aqft.py b/pennylane/templates/subroutines/aqft.py index 190b0b30b61..04144ca15f6 100644 --- a/pennylane/templates/subroutines/aqft.py +++ b/pennylane/templates/subroutines/aqft.py @@ -84,7 +84,7 @@ def circ(): The resulting circuit is: - >>> print(qml.draw(circ, expansion_strategy='device')()) + >>> print(qml.draw(circ, level='device')()) UserWarning: order=0, applying Hadamard transform warnings.warn("order=0, applying Hadamard transform") 0: ──H─╭SWAP────────────── β•­Probs 1: ──H─│─────╭SWAP──────── β”œProbs @@ -105,7 +105,7 @@ def circ(): The resulting circuit is: - >>> print(qml.draw(circ, expansion_strategy='device')()) + >>> print(qml.draw(circ, level='device')()) 0: ──H─╭RΟ•(1.57)─╭RΟ•(0.79)────────────────────────────────────────╭SWAP──────── β•­Probs 1: ────╰●────────│──────────H─╭RΟ•(1.57)─╭RΟ•(0.79)─────────────────│─────╭SWAP── β”œProbs 2: ──────────────╰●───────────╰●────────│──────────H─╭RΟ•(1.57)────│─────╰SWAP── β”œProbs diff --git a/pennylane/templates/subroutines/interferometer.py b/pennylane/templates/subroutines/interferometer.py index d009b7ccb95..1b58774ecdc 100644 --- a/pennylane/templates/subroutines/interferometer.py +++ b/pennylane/templates/subroutines/interferometer.py @@ -124,7 +124,7 @@ def circuit(params): Using these random parameters, the resulting circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")(params)) + >>> print(qml.draw(circuit, level="device")(params)) 0: ─╭BS(0.97,0.09)────────────────╭BS(0.89,0.33)──R(0.83)───────────────── 1: ─╰BS(0.97,0.09)─╭BS(0.94,0.05)─╰BS(0.89,0.33)─╭BS(0.92,0.27)──R(0.36)── 2: ─╭BS(0.78,0.20)─╰BS(0.94,0.05)─╭BS(0.60,0.39)─╰BS(0.92,0.27)──R(0.28)── @@ -144,7 +144,7 @@ def circuit(params): for shape in shapes: params.append(np.random.random(shape)) - print(qml.draw(circuit, expansion_strategy="device")(params)) + print(qml.draw(circuit, level="device")(params)) .. code-block:: diff --git a/pennylane/templates/subroutines/permute.py b/pennylane/templates/subroutines/permute.py index 28a3fd43e65..7e0327f3d8c 100644 --- a/pennylane/templates/subroutines/permute.py +++ b/pennylane/templates/subroutines/permute.py @@ -67,7 +67,7 @@ def apply_perm(): qml.Permute([3, 2, 0, 1], dev.wires) return qml.expval(qml.Z(0)) - >>> print(qml.draw(apply_perm, expansion_strategy="device")()) + >>> print(qml.draw(apply_perm, level="device")()) 0: ─╭SWAP────────────── 1: ─│─────╭SWAP──────── 2: ─│─────╰SWAP─╭SWAP── @@ -108,7 +108,7 @@ def circuit(): The permuted circuit is: - >>> print(qml.draw(circuit, expansion_strategy="device")()) + >>> print(qml.draw(circuit, level="device")()) 3: ─╭SWAP────────────── 2: ─│─────╭SWAP──────── 0: ─│─────│─────╭SWAP── @@ -131,7 +131,7 @@ def circuit(): will permute only the second, third, and fifth wires as follows: - >>> print(qml.draw(circuit, expansion_strategy="device", show_all_wires=True)()) + >>> print(qml.draw(circuit, level="device", show_all_wires=True)()) 3: ────────────── 2: ─╭SWAP──────── a: ─│──────────── diff --git a/pennylane/templates/subroutines/prepselprep.py b/pennylane/templates/subroutines/prepselprep.py index f6c02265a1c..53df7f96cfc 100644 --- a/pennylane/templates/subroutines/prepselprep.py +++ b/pennylane/templates/subroutines/prepselprep.py @@ -75,9 +75,10 @@ class PrepSelPrep(Operation): """ def __init__(self, lcu, control=None, id=None): + coeffs, ops = lcu.terms() control = qml.wires.Wires(control) - self.hyperparameters["lcu"] = lcu + self.hyperparameters["lcu"] = qml.ops.LinearCombination(coeffs, ops) self.hyperparameters["coeffs"] = coeffs self.hyperparameters["ops"] = ops self.hyperparameters["control"] = control @@ -95,14 +96,11 @@ def __init__(self, lcu, control=None, id=None): super().__init__(*self.data, wires=all_wires, id=id) def _flatten(self): - return tuple(self.lcu), (self.control) + return (self.lcu,), (self.control,) @classmethod def _unflatten(cls, data, metadata) -> "PrepSelPrep": - coeffs = [term.terms()[0][0] for term in data] - ops = [term.terms()[1][0] for term in data] - lcu = qml.ops.LinearCombination(coeffs, ops) - return cls(lcu, metadata) + return cls(data[0], metadata[0]) def __repr__(self): return f"PrepSelPrep(coeffs={tuple(self.coeffs)}, ops={tuple(self.ops)}, control={self.control})" diff --git a/pennylane/templates/subroutines/select.py b/pennylane/templates/subroutines/select.py index 357cf0113b0..4bab8074b64 100644 --- a/pennylane/templates/subroutines/select.py +++ b/pennylane/templates/subroutines/select.py @@ -56,7 +56,7 @@ class Select(Operation): >>> qml.Select(ops, control=[0,1]) >>> return qml.state() ... - >>> print(qml.draw(circuit, expansion_strategy='device')()) + >>> print(qml.draw(circuit, level='device')()) 0: ─╭○─╭○─╭●─╭●───── State 1: β”€β”œβ—‹β”€β”œβ—β”€β”œβ—‹β”€β”œβ—β”€β”€β”€β”€β”€ State 2: ─╰X─│──╰Yβ”€β”œSWAP── State diff --git a/pennylane/templates/subroutines/trotter.py b/pennylane/templates/subroutines/trotter.py index 64c26dc4768..ae5aa6c5031 100644 --- a/pennylane/templates/subroutines/trotter.py +++ b/pennylane/templates/subroutines/trotter.py @@ -438,10 +438,9 @@ def compute_decomposition(*args, **kwargs): ops = kwargs["base"].operands decomp = _recursive_expression(time / n, order, ops)[::-1] * n - unique_decomp = [copy.copy(op) for op in decomp] if qml.QueuingManager.recording(): - for op in unique_decomp: # apply operators in reverse order of expression + for op in decomp: # apply operators in reverse order of expression qml.apply(op) - return unique_decomp + return decomp diff --git a/pennylane/templates/swapnetworks/ccl2.py b/pennylane/templates/swapnetworks/ccl2.py index 378ff9e389c..9595e577e0c 100644 --- a/pennylane/templates/swapnetworks/ccl2.py +++ b/pennylane/templates/swapnetworks/ccl2.py @@ -54,7 +54,7 @@ class TwoLocalSwapNetwork(Operation): ... def swap_network_circuit(): ... qml.templates.TwoLocalSwapNetwork(dev.wires, acquaintances, fermionic=True, shift=False) ... return qml.state() - >>> qml.draw(swap_network_circuit, expansion_strategy='device')() + >>> qml.draw(swap_network_circuit, level='device')() 0: ─╭●─╭fSWAP(3.14)─────────────────╭●─╭fSWAP(3.14)─────────────────╭●─╭fSWAP(3.14)── State 1: ─╰X─╰fSWAP(3.14)─╭●─╭fSWAP(3.14)─╰X─╰fSWAP(3.14)─╭●─╭fSWAP(3.14)─╰X─╰fSWAP(3.14)── State 2: ─╭●─╭fSWAP(3.14)─╰X─╰fSWAP(3.14)─╭●─╭fSWAP(3.14)─╰X─╰fSWAP(3.14)─╭●─╭fSWAP(3.14)── State @@ -79,7 +79,7 @@ class TwoLocalSwapNetwork(Operation): ... def swap_network_circuit(): ... qml.templates.TwoLocalSwapNetwork(dev.wires, acquaintances, weights, fermionic=False) ... return qml.state() - >>> qml.draw(swap_network_circuit, expansion_strategy='device')() + >>> qml.draw(swap_network_circuit, level='device')() 0: ─╭●────────╭SWAP─────────────────╭●────────╭SWAP─────────────────╭●────────╭SWAP── State 1: ─╰RY(0.20)─╰SWAP─╭●────────╭SWAP─╰RY(0.09)─╰SWAP─╭●────────╭SWAP─╰RY(0.62)─╰SWAP── State 2: ─╭●────────╭SWAP─╰RY(0.68)─╰SWAP─╭●────────╭SWAP─╰RY(0.34)─╰SWAP─╭●────────╭SWAP── State diff --git a/pennylane/templates/tensornetworks/mera.py b/pennylane/templates/tensornetworks/mera.py index 5a4b6a1ee35..1adfcad2322 100644 --- a/pennylane/templates/tensornetworks/mera.py +++ b/pennylane/templates/tensornetworks/mera.py @@ -16,7 +16,7 @@ """ # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access import warnings -from typing import Callable +from collections.abc import Callable import numpy as np @@ -160,7 +160,7 @@ def circuit(template_weights): It may be necessary to reorder the wires to see the MERA architecture clearly: - >>> print(qml.draw(circuit, expansion_strategy='device', wire_order=[2,0,1,3])(template_weights)) + >>> print(qml.draw(circuit, level='device', wire_order=[2,0,1,3])(template_weights)) 2: ───────────────╭●──RY(0.10)──╭X──RY(-0.30)──────────────── 0: ─╭X──RY(-0.30)─│─────────────╰●──RY(0.10)──╭●──RY(0.10)─── 1: ─╰●──RY(0.10)──│─────────────╭X──RY(-0.30)─╰X──RY(-0.30)── diff --git a/pennylane/templates/tensornetworks/mps.py b/pennylane/templates/tensornetworks/mps.py index 922c5d531fa..4f7922e6836 100644 --- a/pennylane/templates/tensornetworks/mps.py +++ b/pennylane/templates/tensornetworks/mps.py @@ -122,7 +122,7 @@ def circuit(template_weights): qml.MPS(range(n_wires),n_block_wires,block, n_params_block, template_weights) return qml.expval(qml.Z(n_wires-1)) - >>> print(qml.draw(circuit, expansion_strategy='device')(template_weights)) + >>> print(qml.draw(circuit, level='device')(template_weights)) 0: ─╭●──RY(0.10)─────────────────────────────── 1: ─╰X──RY(-0.30)─╭●──RY(0.10)───────────────── 2: ───────────────╰X──RY(-0.30)─╭●──RY(0.10)─── @@ -148,7 +148,7 @@ def circuit(): qml.MPS(range(n_wires),n_block_wires, block, n_params_block, offset = 1) return qml.state() - >>> print(qml.draw(circuit, expansion_strategy='device')()) + >>> print(qml.draw(circuit, level='device')()) 0: ─╭●────────────── State 1: β”€β”œβ—β”€β•­β—β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ State 2: β”€β”œβ—β”€β”œβ—β”€β•­β—β”€β”€β”€β”€β”€β”€β”€β”€ State diff --git a/pennylane/templates/tensornetworks/ttn.py b/pennylane/templates/tensornetworks/ttn.py index 7bd88ce2e9d..40029f4284f 100644 --- a/pennylane/templates/tensornetworks/ttn.py +++ b/pennylane/templates/tensornetworks/ttn.py @@ -131,7 +131,7 @@ def circuit(template_weights): qml.TTN(range(n_wires),n_block_wires,block, n_params_block, template_weights) return qml.expval(qml.Z(n_wires-1)) - >>> print(qml.draw(circuit, expansion_strategy='device')(template_weights)) + >>> print(qml.draw(circuit, level='device')(template_weights)) 0: ─╭●──RY(0.10)───────────────── 1: ─╰X──RY(-0.30)─╭●──RY(0.10)─── 2: ─╭●──RY(0.10)──│────────────── diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py index 83dcb1ae569..b232176ec25 100644 --- a/pennylane/transforms/__init__.py +++ b/pennylane/transforms/__init__.py @@ -93,6 +93,7 @@ ~transforms.fold_global ~transforms.poly_extrapolate ~transforms.richardson_extrapolate + ~transforms.exponential_extrapolate Other transforms ~~~~~~~~~~~~~~~~ @@ -199,10 +200,10 @@ .. code-block:: python - from typing import Sequence, Callable - from pennylane.tape import QuantumTape + from pennylane.tape import QuantumTape, QuantumTapeBatch + from pennylane.typing import PostprocessingFn - def remove_rx(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): + def remove_rx(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: operations = filter(lambda op: op.name != "RX", tape.operations) new_tape = type(tape)(operations, tape.measurements, shots=tape.shots) @@ -226,11 +227,11 @@ def null_postprocessing(results): .. code-block:: python - from typing import Sequence, Callable - from pennylane.tape import QuantumTape + from pennylane.tape import QuantumTape, QuantumTapeBatch + from pennylane.typing import PostprocessingFn @qml.transform - def sum_circuit_and_adjoint(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): + def sum_circuit_and_adjoint(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: operations = [qml.adjoint(op) for op in tape.operation] new_tape = type(tape)(operations, tape.measurements, shots=tape.shots) @@ -322,7 +323,13 @@ def circuit(params): from .split_to_single_terms import split_to_single_terms from .insert_ops import insert -from .mitigate import mitigate_with_zne, fold_global, poly_extrapolate, richardson_extrapolate +from .mitigate import ( + mitigate_with_zne, + fold_global, + poly_extrapolate, + richardson_extrapolate, + exponential_extrapolate, +) from .optimization import ( cancel_inverses, commute_controlled, diff --git a/pennylane/transforms/batch_input.py b/pennylane/transforms/batch_input.py index 9d53c2fe97f..538b61bc250 100644 --- a/pennylane/transforms/batch_input.py +++ b/pennylane/transforms/batch_input.py @@ -13,21 +13,23 @@ """ Batch transformation for multiple (non-trainable) input examples following issue #2037 """ -from typing import Callable, Sequence, Union +from collections.abc import Sequence +from typing import Union import numpy as np import pennylane as qml -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.batch_params import _nested_stack, _split_operations from pennylane.transforms.core import transform +from pennylane.typing import PostprocessingFn @transform def batch_input( tape: QuantumTape, argnum: Union[Sequence[int], int], -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """ Transform a circuit to support an initial batch dimension for gate inputs. diff --git a/pennylane/transforms/batch_params.py b/pennylane/transforms/batch_params.py index 9aa0598a0dd..aefe5bf1ae5 100644 --- a/pennylane/transforms/batch_params.py +++ b/pennylane/transforms/batch_params.py @@ -15,9 +15,10 @@ Contains the batch dimension transform. """ # pylint: disable=import-outside-toplevel -from typing import Callable, Sequence import pennylane as qml +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .core import transform @@ -82,7 +83,7 @@ def _split_operations(ops, params, split_indices, num_tapes): @transform def batch_params( tape: qml.tape.QuantumTape, all_operations=False -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transform a QNode to support an initial batch dimension for operation parameters. diff --git a/pennylane/transforms/broadcast_expand.py b/pennylane/transforms/broadcast_expand.py index cdd4c5a8192..68f9a2d1f98 100644 --- a/pennylane/transforms/broadcast_expand.py +++ b/pennylane/transforms/broadcast_expand.py @@ -13,10 +13,11 @@ # limitations under the License. """This module contains the tape expansion function for expanding a broadcasted tape into multiple tapes.""" -from typing import Callable, Sequence import pennylane as qml from pennylane.measurements import MidMeasureMP, SampleMP +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .core import transform @@ -48,7 +49,7 @@ def _split_operations(ops, num_tapes): @transform -def broadcast_expand(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable): +def broadcast_expand(tape: qml.tape.QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Expand a broadcasted tape into multiple tapes and a function that stacks and squeezes the results. diff --git a/pennylane/transforms/commutation_dag.py b/pennylane/transforms/commutation_dag.py index 557676c45f5..c0b62c743a6 100644 --- a/pennylane/transforms/commutation_dag.py +++ b/pennylane/transforms/commutation_dag.py @@ -17,19 +17,19 @@ import heapq from collections import OrderedDict from functools import partial -from typing import Callable, Sequence import networkx as nx from networkx.drawing.nx_pydot import to_pydot import pennylane as qml -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires @partial(transform, is_informative=True) -def commutation_dag(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def commutation_dag(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Construct the pairwise-commutation DAG (directed acyclic graph) representation of a quantum circuit. In the DAG, each node represents a quantum operation, and edges represent diff --git a/pennylane/transforms/compile.py b/pennylane/transforms/compile.py index 50bb7104989..80deaf295aa 100644 --- a/pennylane/transforms/compile.py +++ b/pennylane/transforms/compile.py @@ -14,12 +14,11 @@ """Code for the high-level quantum function transform that executes compilation.""" # pylint: disable=too-many-branches from functools import partial -from typing import Callable, Sequence import pennylane as qml from pennylane.ops import __all__ as all_ops from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.core import TransformDispatcher, transform from pennylane.transforms.optimization import ( cancel_inverses, @@ -27,6 +26,7 @@ merge_rotations, remove_barrier, ) +from pennylane.typing import PostprocessingFn default_pipeline = [commute_controlled, cancel_inverses, merge_rotations, remove_barrier] @@ -34,7 +34,7 @@ @transform def compile( tape: QuantumTape, pipeline=None, basis_set=None, num_passes=1, expand_depth=5 -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Compile a circuit by applying a series of transforms to a quantum function. The default set of transforms includes (in order): diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py index 1167fec69c4..83a3343cc21 100644 --- a/pennylane/transforms/convert_to_numpy_parameters.py +++ b/pennylane/transforms/convert_to_numpy_parameters.py @@ -15,12 +15,12 @@ This file contains preprocessings steps that may be called internally during execution. """ -from typing import Callable, Sequence, Tuple import pennylane as qml from pennylane import math -from pennylane.tape import QuantumScript +from pennylane.tape import QuantumScript, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn # pylint: disable=no-member @@ -48,7 +48,7 @@ def _convert_measurement_to_numpy_data( # pylint: disable=protected-access @transform -def convert_to_numpy_parameters(tape: QuantumScript) -> Tuple[Sequence[QuantumScript], Callable]: +def convert_to_numpy_parameters(tape: QuantumScript) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transforms a circuit to one with purely numpy parameters. Args: diff --git a/pennylane/transforms/core/transform.py b/pennylane/transforms/core/transform.py index fc4c3b72198..6c1a632b44d 100644 --- a/pennylane/transforms/core/transform.py +++ b/pennylane/transforms/core/transform.py @@ -45,7 +45,7 @@ def transform( returns a sequence of :class:`~.QuantumTape` and a processing function. * The transform must have the following structure (type hinting is optional): ``my_quantum_transform(tape: - qml.tape.QuantumTape, ...) -> ( Sequence[qml.tape.QuantumTape], Callable)`` + qml.tape.QuantumTape, ...) -> tuple[qml.tape.QuantumTapeBatch, qml.typing.PostprocessingFn]`` Keyword Args: expand_transform=None (Optional[Callable]): An optional expand transform is applied directly before the input @@ -72,9 +72,10 @@ def transform( .. code-block:: python - from typing import Sequence, Callable + from pennylane.tape import QuantumTapeBatch + from pennylane.typing import PostprocessingFn - def my_quantum_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable): + def my_quantum_transform(tape: qml.tape.QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: tape1 = tape tape2 = tape.copy() @@ -146,13 +147,13 @@ def qnode_circuit(a): operations = [qml.Hadamard(0), qml.RX(0.2, 0), qml.RX(0.6, 0), qml.CNOT((0, 1))] tape = qml.tape.QuantumTape(operations, measurement) - batch1, function1 = qml.transforms.hamiltonian_expand(tape) + batch1, function1 = qml.transforms.split_non_commuting(tape) batch2, function2 = qml.transforms.merge_rotations(batch1) dev = qml.device("default.qubit", wires=3) result = dev.execute(batch2) - The first ``hamiltonian_expand`` transform splits the original tape, returning a batch of tapes ``batch1`` and a processing function ``function1``. + The first ``split_non_commuting`` transform splits the original tape, returning a batch of tapes ``batch1`` and a processing function ``function1``. The second ``merge_rotations`` transform is applied to the batch of tapes returned by the first transform. It returns a new batch of tapes ``batch2``, each of which has been transformed by the second transform, and a processing function ``function2``. diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py index 7591c9d2222..08457859baa 100644 --- a/pennylane/transforms/core/transform_dispatcher.py +++ b/pennylane/transforms/core/transform_dispatcher.py @@ -19,7 +19,7 @@ import os import types import warnings -from typing import Sequence +from collections.abc import Sequence import pennylane as qml from pennylane.typing import ResultBatch diff --git a/pennylane/transforms/core/transform_program.py b/pennylane/transforms/core/transform_program.py index 14d492f704e..b91576bfc59 100644 --- a/pennylane/transforms/core/transform_program.py +++ b/pennylane/transforms/core/transform_program.py @@ -14,21 +14,19 @@ """ This module contains the ``TransformProgram`` class. """ +from collections.abc import Sequence from functools import partial -from typing import Callable, List, Optional, Sequence, Tuple, Union +from typing import Optional, Union import pennylane as qml -from pennylane.tape import QuantumTape -from pennylane.typing import Result, ResultBatch +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import BatchPostprocessingFn, PostprocessingFn, ResultBatch from .transform_dispatcher import TransformContainer, TransformDispatcher, TransformError -PostProcessingFn = Callable[[ResultBatch], Result] -BatchPostProcessingFn = Callable[[ResultBatch], ResultBatch] - def _batch_postprocessing( - results: ResultBatch, individual_fns: List[PostProcessingFn], slices: List[slice] + results: ResultBatch, individual_fns: list[PostprocessingFn], slices: list[slice] ) -> ResultBatch: """Broadcast individual post processing functions onto their respective tapes. @@ -58,7 +56,7 @@ def _batch_postprocessing( def _apply_postprocessing_stack( results: ResultBatch, - postprocessing_stack: List[BatchPostProcessingFn], + postprocessing_stack: list[BatchPostprocessingFn], ) -> ResultBatch: """Applies the postprocessing and cotransform postprocessing functions in a Last-In-First-Out LIFO manner. @@ -141,7 +139,7 @@ class TransformProgram: True >>> qml.compile in program True - >>> qml.transforms.hamiltonian_expand in program + >>> qml.transforms.split_non_commuting in program False >>> program + program TransformProgram(compile, cancel_inverses, compile, cancel_inverses) @@ -491,7 +489,7 @@ def _set_all_argnums(self, qnode, args, kwargs, argnums): qnode.construct(args, kwargs) - def __call__(self, tapes: Tuple[QuantumTape]) -> Tuple[ResultBatch, BatchPostProcessingFn]: + def __call__(self, tapes: QuantumTapeBatch) -> tuple[QuantumTapeBatch, BatchPostprocessingFn]: if not self: return tapes, null_postprocessing diff --git a/pennylane/transforms/decompositions/clifford_t_transform.py b/pennylane/transforms/decompositions/clifford_t_transform.py index 92698752e9e..6b7cb16cf4b 100644 --- a/pennylane/transforms/decompositions/clifford_t_transform.py +++ b/pennylane/transforms/decompositions/clifford_t_transform.py @@ -16,13 +16,12 @@ import math import warnings from itertools import product -from typing import Callable, Sequence import pennylane as qml from pennylane.ops import Adjoint from pennylane.ops.op_math.decompositions.solovay_kitaev import sk_decomposition from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms.core import transform from pennylane.transforms.optimization import ( cancel_inverses, @@ -31,6 +30,7 @@ remove_barrier, ) from pennylane.transforms.optimization.optimization_utils import _fuse_global_phases, find_next_gate +from pennylane.typing import PostprocessingFn # Single qubits Clifford+T gates in PL _CLIFFORD_T_ONE_GATES = [ @@ -83,7 +83,7 @@ def _check_clifford_op(op, use_decomposition=False): # Check if matrix can be calculated for the operator if (not op.has_matrix and not use_decomposition) or ( - use_decomposition and not op.expand().wires + use_decomposition and not qml.tape.QuantumScript(op.decomposition()).wires ): return False @@ -314,7 +314,7 @@ def clifford_t_decomposition( max_expansion=6, method="sk", **method_kwargs, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Decomposes a circuit into the Clifford+T basis. This method first decomposes the gate operations to a basis comprised of Clifford, :class:`~.T`, :class:`~.RZ` and diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py index bf1bc5f568c..323fd676050 100644 --- a/pennylane/transforms/defer_measurements.py +++ b/pennylane/transforms/defer_measurements.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. """Code for the tape transform implementing the deferred measurement principle.""" -from typing import Callable, Sequence import pennylane as qml from pennylane.measurements import CountsMP, MeasurementValue, MidMeasureMP, ProbabilityMP, SampleMP from pennylane.ops.op_math import ctrl from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires # pylint: disable=too-many-branches, protected-access, too-many-statements @@ -104,7 +104,7 @@ def null_postprocessing(results): @transform def defer_measurements( tape: QuantumTape, reduce_postselected: bool = True, **kwargs -) -> tuple[Sequence[QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Quantum function transform that substitutes operations conditioned on measurement outcomes to controlled operations. diff --git a/pennylane/transforms/dynamic_one_shot.py b/pennylane/transforms/dynamic_one_shot.py index 591553b3021..2eb326d29d1 100644 --- a/pennylane/transforms/dynamic_one_shot.py +++ b/pennylane/transforms/dynamic_one_shot.py @@ -19,7 +19,7 @@ # pylint: disable=import-outside-toplevel from collections import Counter -from typing import Callable, Sequence +from collections.abc import Sequence import numpy as np @@ -33,7 +33,8 @@ SampleMP, VarianceMP, ) -from pennylane.typing import TensorLike +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn, TensorLike from .core import transform @@ -56,7 +57,7 @@ def null_postprocessing(results): @transform def dynamic_one_shot( tape: qml.tape.QuantumTape, **kwargs -) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transform a QNode to into several one-shot tapes to support dynamic circuit execution. Args: diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py index 76b693f6175..62370e65dd2 100644 --- a/pennylane/transforms/hamiltonian_expand.py +++ b/pennylane/transforms/hamiltonian_expand.py @@ -14,16 +14,18 @@ """ Contains the hamiltonian expand tape transform """ + # pylint: disable=protected-access +import warnings +from collections.abc import Sequence from functools import partial -from typing import Callable, List, Sequence, Tuple import pennylane as qml from pennylane.measurements import ExpectationMP, MeasurementProcess, Shots from pennylane.ops import Prod, SProd, Sum -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform -from pennylane.typing import ResultBatch +from pennylane.typing import PostprocessingFn, ResultBatch def grouping_processing_fn(res_groupings, coeff_groupings, batch_size, offset): @@ -137,7 +139,9 @@ def _naive_hamiltonian_expand(tape): @transform -def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], Callable): +def hamiltonian_expand( + tape: QuantumTape, group: bool = True +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r""" Splits a tape measuring a Hamiltonian expectation into mutliple tapes of Pauli expectations, and provides a function to recombine the results. @@ -151,6 +155,10 @@ def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[Quant Returns: qnode (QNode) or tuple[List[QuantumTape], function]: The transformed circuit as described in :func:`qml.transform `. + .. warning:: + This function is deprecated and will be removed in version 0.39. + Instead, use :func:`~.transforms.split_non_commuting`. + **Example** Given a Hamiltonian, @@ -223,6 +231,12 @@ def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[Quant 2 """ + warnings.warn( + "qml.transforms.hamiltonian_expand is deprecated and will be removed in version 0.39. " + "Instead, use qml.transforms.split_non_commuting, which can handle the same measurement type.", + qml.PennyLaneDeprecationWarning, + ) + if ( len(tape.measurements) != 1 or not hasattr(tape.measurements[0].obs, "grouping_indices") @@ -244,8 +258,8 @@ def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[Quant def _group_measurements( - measurements: Sequence[MeasurementProcess], indices_and_coeffs: List[List[Tuple[int, float]]] -) -> (List[List[MeasurementProcess]], List[List[Tuple[int, int, float]]]): + measurements: Sequence[MeasurementProcess], indices_and_coeffs: list[list[tuple[int, float]]] +) -> tuple[list[list[MeasurementProcess]], list[list[tuple[int, int, float]]]]: """Groups measurements that does not have overlapping wires. Returns: @@ -318,10 +332,10 @@ def _group_measurements( def _sum_expand_processing_fn_grouping( res: ResultBatch, - group_sizes: List[int], + group_sizes: list[int], shots: Shots, - indices_and_coeffs: List[List[Tuple[int, int, float]]], - offsets: List[int], + indices_and_coeffs: list[list[tuple[int, int, float]]], + offsets: list[int], ): """The processing function for sum_expand with grouping.""" @@ -348,8 +362,8 @@ def _sum_expand_processing_fn_grouping( def _sum_expand_processing_fn( res: ResultBatch, shots: Shots, - indices_and_coeffs: List[List[Tuple[int, float]]], - offsets: List[int], + indices_and_coeffs: list[list[tuple[int, float]]], + offsets: list[int], ): """The processing function for sum_expand without grouping.""" @@ -369,7 +383,7 @@ def _sum_expand_processing_fn( @transform -def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], Callable): +def sum_expand(tape: QuantumTape, group: bool = True) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Splits a quantum tape measuring a Sum expectation into multiple tapes of summand expectations, and provides a function to recombine the results. @@ -384,6 +398,10 @@ def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], quantum tapes to be evaluated, and a function to be applied to these tape executions to compute the expectation value. + .. warning:: + This function is deprecated and will be removed in version 0.39. + Instead, use :func:`~.transforms.split_non_commuting`. + **Example** Given a Sum operator, @@ -458,6 +476,12 @@ def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], """ + warnings.warn( + "qml.transforms.sum_expand is deprecated and will be removed in version 0.39. " + "Instead, use qml.transforms.split_non_commuting, which can handle the same measurement type.", + qml.PennyLaneDeprecationWarning, + ) + # The dictionary of all unique single-term observable measurements, and their indices # within the list of all single-term observable measurements. single_term_obs_measurements = {} diff --git a/pennylane/transforms/insert_ops.py b/pennylane/transforms/insert_ops.py index 17a81beeec8..adff13e1494 100644 --- a/pennylane/transforms/insert_ops.py +++ b/pennylane/transforms/insert_ops.py @@ -14,14 +14,16 @@ """ Provides transforms for inserting operations into quantum circuits. """ +from collections.abc import Sequence from types import FunctionType -from typing import Callable, Sequence, Type, Union +from typing import Type, Union import pennylane as qml from pennylane.operation import Operation from pennylane.ops.op_math import Adjoint -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn # pylint: disable=too-many-branches @@ -56,7 +58,7 @@ def insert( op_args: Union[tuple, float], position: Union[str, list, Type[Operation]] = "all", before: bool = False, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Insert an operation into specified points in an input circuit. Circuits passed through this transform will be updated to have the operation, specified by the diff --git a/pennylane/transforms/mitigate.py b/pennylane/transforms/mitigate.py index 572c904bc44..f4f7f32856f 100644 --- a/pennylane/transforms/mitigate.py +++ b/pennylane/transforms/mitigate.py @@ -12,19 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. """Provides transforms for mitigating quantum circuits.""" +from collections.abc import Sequence from copy import copy -from typing import Any, Callable, Dict, Optional, Sequence +from typing import Any, Optional import pennylane as qml from pennylane import adjoint, apply from pennylane.math import mean, round, shape from pennylane.queuing import AnnotatedQueue -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn @transform -def fold_global(tape: QuantumTape, scale_factor) -> (Sequence[QuantumTape], Callable): +def fold_global(tape: QuantumTape, scale_factor) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Differentiable circuit folding of the global unitary ``circuit``. For a unitary circuit :math:`U = L_d .. L_1`, where :math:`L_i` can be either a gate or layer, ``fold_global`` constructs @@ -318,6 +320,48 @@ def richardson_extrapolate(x, y): return poly_extrapolate(x, y, len(x) - 1) +def exponential_extrapolate(x, y, asymptote=None, eps=1.0e-6): + r"""Extrapolate to the zero-noise limit using an exponential model (:math:`Ae^{Bx} + C`). This + is done by linearizing the data using a logarithm, whereupon a linear fit is performed. Once + the model parameters are found, they are transformed back to exponential parameters. + + Args: + x (Array): Data in x axis. + y (Array): Data in y axis such that :math:`y = f(x)`. + asymptote (float): Infinite noise limit expected for your circuit of interest (:math:`C` + in the equation above). Defaults to 0 in the case an asymptote is not supplied. + eps (float): Epsilon to regularize :math:`\log(y - C)` when the argument is to close to + zero or negative. + + Returns: + float: Extrapolated value at f(0). + + .. seealso:: :func:`~.pennylane.transforms.richardson_extrapolate`, :func:`~.pennylane.transforms.mitigate_with_zne`. + + **Example:** + + >>> np.random.seed(0) + >>> x = np.linspace(1, 10, 5) + >>> y = np.exp(-x) + np.random.normal(scale=0.1, size=len(x)) + >>> qml.transforms.exponential_extrapolate(x, y) + 0.23365009000522544 + """ + y = qml.math.stack(y) + slope, y_intercept = _polyfit(x, y, 1) + if asymptote is None: + sign = qml.math.sign(-slope) + asymptote = 0.0 + else: + sign = qml.math.sign(-(asymptote - y_intercept)) + + y_shifted = sign * (y - asymptote) + y_shifted = qml.math.where(y_shifted < eps, eps, y_shifted) + y_scaled = qml.math.log(y_shifted) + + zne_unscaled = poly_extrapolate(x, y_scaled, 1) + return sign * qml.math.exp(zne_unscaled) + asymptote + + # pylint: disable=too-many-arguments, protected-access @transform def mitigate_with_zne( @@ -325,10 +369,10 @@ def mitigate_with_zne( scale_factors: Sequence[float], folding: callable, extrapolate: callable, - folding_kwargs: Optional[Dict[str, Any]] = None, - extrapolate_kwargs: Optional[Dict[str, Any]] = None, + folding_kwargs: Optional[dict[str, Any]] = None, + extrapolate_kwargs: Optional[dict[str, Any]] = None, reps_per_factor=1, -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Mitigate an input circuit using zero-noise extrapolation. Error mitigation is a precursor to error correction and is compatible with near-term quantum diff --git a/pennylane/transforms/optimization/cancel_inverses.py b/pennylane/transforms/optimization/cancel_inverses.py index b4d259ae3dc..fc9f4854d33 100644 --- a/pennylane/transforms/optimization/cancel_inverses.py +++ b/pennylane/transforms/optimization/cancel_inverses.py @@ -13,7 +13,6 @@ # limitations under the License. """Transform for cancelling adjacent inverse gates in quantum circuits.""" # pylint: disable=too-many-branches -from typing import Callable, Sequence from pennylane.ops.op_math import Adjoint from pennylane.ops.qubit.attributes import ( @@ -21,8 +20,9 @@ symmetric_over_all_wires, symmetric_over_control_wires, ) -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires from .optimization_utils import find_next_gate @@ -64,7 +64,7 @@ def _are_inverses(op1, op2): @transform -def cancel_inverses(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def cancel_inverses(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Quantum function transform to remove any operations that are applied next to their (self-)inverses or adjoint. diff --git a/pennylane/transforms/optimization/commute_controlled.py b/pennylane/transforms/optimization/commute_controlled.py index 9314b051bfe..896f4824f9c 100644 --- a/pennylane/transforms/optimization/commute_controlled.py +++ b/pennylane/transforms/optimization/commute_controlled.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. """Transforms for pushing commuting gates through targets/control qubits.""" -from typing import Callable, Sequence -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires from .optimization_utils import find_next_gate @@ -153,7 +153,9 @@ def _commute_controlled_left(op_list): @transform -def commute_controlled(tape: QuantumTape, direction="right") -> (Sequence[QuantumTape], Callable): +def commute_controlled( + tape: QuantumTape, direction="right" +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Quantum transform to move commuting gates past control and target qubits of controlled operations. Args: diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py index 34ad6c57009..01567048943 100644 --- a/pennylane/transforms/optimization/merge_amplitude_embedding.py +++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py @@ -12,18 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. """Transform for merging AmplitudeEmbedding gates in a quantum circuit.""" -from typing import Callable, Sequence from pennylane import AmplitudeEmbedding from pennylane._device import DeviceError from pennylane.math import flatten, reshape from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn @transform -def merge_amplitude_embedding(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def merge_amplitude_embedding(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Quantum function transform to combine amplitude embedding templates that act on different qubits. Args: diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py index 2c8178c1997..38cfaefd7ba 100644 --- a/pennylane/transforms/optimization/merge_rotations.py +++ b/pennylane/transforms/optimization/merge_rotations.py @@ -13,15 +13,15 @@ # limitations under the License. """Transform for merging adjacent rotations of the same type in a quantum circuit.""" # pylint: disable=too-many-branches -from typing import Callable, Sequence import pennylane as qml from pennylane.math import allclose, cast_like, get_interface, is_abstract, stack, zeros from pennylane.ops.op_math import Adjoint from pennylane.ops.qubit.attributes import composable_rotations from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from .optimization_utils import find_next_gate, fuse_rot_angles @@ -29,7 +29,7 @@ @transform def merge_rotations( tape: QuantumTape, atol=1e-8, include_gates=None -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Quantum transform to combine rotation gates of the same type that act sequentially. If the combination of two rotation produces an angle that is close to 0, diff --git a/pennylane/transforms/optimization/pattern_matching.py b/pennylane/transforms/optimization/pattern_matching.py index c14a2e9c203..0eb34b1a98e 100644 --- a/pennylane/transforms/optimization/pattern_matching.py +++ b/pennylane/transforms/optimization/pattern_matching.py @@ -17,16 +17,16 @@ import copy import itertools from collections import OrderedDict -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane import adjoint from pennylane.ops.qubit.attributes import symmetric_over_all_wires -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch from pennylane.transforms import transform from pennylane.transforms.commutation_dag import commutation_dag +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires @@ -34,7 +34,7 @@ @transform def pattern_matching_optimization( tape: QuantumTape, pattern_tapes, custom_quantum_cost=None -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Quantum function transform to optimize a circuit given a list of patterns (templates). Args: diff --git a/pennylane/transforms/optimization/remove_barrier.py b/pennylane/transforms/optimization/remove_barrier.py index fdb847caab6..e903d0535ee 100644 --- a/pennylane/transforms/optimization/remove_barrier.py +++ b/pennylane/transforms/optimization/remove_barrier.py @@ -13,14 +13,14 @@ # limitations under the License. """Transform for removing the Barrier gate from quantum circuits.""" # pylint: disable=too-many-branches -from typing import Callable, Sequence -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn @transform -def remove_barrier(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def remove_barrier(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Quantum transform to remove Barrier gates. Args: diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py index cf1d75459e0..9bb56dbd500 100644 --- a/pennylane/transforms/optimization/single_qubit_fusion.py +++ b/pennylane/transforms/optimization/single_qubit_fusion.py @@ -13,13 +13,13 @@ # limitations under the License. """Transform for fusing sequences of single-qubit gates.""" # pylint: disable=too-many-branches -from typing import Callable, Sequence from pennylane.math import allclose, is_abstract, stack from pennylane.ops.qubit import Rot from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn from .optimization_utils import find_next_gate, fuse_rot_angles @@ -27,7 +27,7 @@ @transform def single_qubit_fusion( tape: QuantumTape, atol=1e-8, exclude_gates=None -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Quantum function transform to fuse together groups of single-qubit operations into a general single-qubit unitary operation (:class:`~.Rot`). diff --git a/pennylane/transforms/optimization/undo_swaps.py b/pennylane/transforms/optimization/undo_swaps.py index 2b64dc17f64..ac5abd9e4ca 100644 --- a/pennylane/transforms/optimization/undo_swaps.py +++ b/pennylane/transforms/optimization/undo_swaps.py @@ -14,10 +14,10 @@ """Transform that eliminates the swap operators by reordering the wires.""" # pylint: disable=too-many-branches -from typing import Callable, Sequence -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn def null_postprocessing(results): @@ -28,7 +28,7 @@ def null_postprocessing(results): @transform -def undo_swaps(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def undo_swaps(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Quantum function transform to remove SWAP gates by running from right to left through the circuit changing the position of the qubits accordingly. diff --git a/pennylane/transforms/qmc.py b/pennylane/transforms/qmc.py index 0f2dd76b34f..e11874a909e 100644 --- a/pennylane/transforms/qmc.py +++ b/pennylane/transforms/qmc.py @@ -15,12 +15,13 @@ Contains the quantum_monte_carlo transform. """ from copy import copy -from typing import Callable, Sequence import pennylane as qml from pennylane import CZ, Hadamard, MultiControlledX, PauliX, adjoint +from pennylane.tape import QuantumTapeBatch from pennylane.templates import QFT from pennylane.transforms.core import transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires @@ -82,7 +83,7 @@ def _apply_controlled_v(target_wire, control_wire): @transform def apply_controlled_Q( tape: qml.tape.QuantumTape, wires, target_wire, control_wire, work_wires -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Applies the transform that performs a controlled version of the :math:`\mathcal{Q}` unitary defined in `this `__ paper. @@ -151,7 +152,7 @@ def apply_controlled_Q( @transform def quantum_monte_carlo( tape: qml.tape.QuantumTape, wires, target_wire, estimation_wires -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Applies the transform `quantum Monte Carlo estimation `__ algorithm. @@ -336,7 +337,7 @@ def qmc(): It is also possible to explore the resources required to perform the quantum Monte Carlo algorithm - >>> qml.specs(qmc, expansion_strategy="device")() + >>> qml.specs(qmc, level="device")() {'resources': Resources( num_wires=12, num_gates=31882, diff --git a/pennylane/transforms/sign_expand/sign_expand.py b/pennylane/transforms/sign_expand/sign_expand.py index 1f0a56efb0e..429bc2a848b 100644 --- a/pennylane/transforms/sign_expand/sign_expand.py +++ b/pennylane/transforms/sign_expand/sign_expand.py @@ -15,12 +15,13 @@ # pylint: disable=protected-access import json from os import path -from typing import Callable, Sequence import numpy as np import pennylane as qml +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn def controlled_pauli_evolution(theta, wires, pauli_word, controls): @@ -199,7 +200,7 @@ def construct_sgn_circuit( # pylint: disable=too-many-arguments @transform def sign_expand( # pylint: disable=too-many-arguments tape: qml.tape.QuantumTape, circuit=False, J=10, delta=0.0, controls=("Hadamard", "Target") -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r""" Splits a tape measuring a (fast-forwardable) Hamiltonian expectation into mutliple tapes of the Xi or sgn decomposition, and provides a function to recombine the results. diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py index a0c1f0f6743..741345f1718 100644 --- a/pennylane/transforms/split_non_commuting.py +++ b/pennylane/transforms/split_non_commuting.py @@ -19,13 +19,14 @@ # pylint: disable=too-many-arguments,too-many-boolean-expressions from functools import partial -from typing import Callable, Dict, List, Optional, Sequence, Tuple +from typing import Optional import pennylane as qml from pennylane.measurements import ExpectationMP, MeasurementProcess, Shots, StateMP from pennylane.ops import Hamiltonian, LinearCombination, Prod, SProd, Sum +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import transform -from pennylane.typing import Result, ResultBatch +from pennylane.typing import PostprocessingFn, Result, ResultBatch def null_postprocessing(results): @@ -39,7 +40,7 @@ def null_postprocessing(results): def split_non_commuting( tape: qml.tape.QuantumScript, grouping_strategy: Optional[str] = "default", -) -> Tuple[Sequence[qml.tape.QuantumTape], Callable]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Splits a circuit into tapes measuring groups of commuting observables. Args: @@ -383,8 +384,8 @@ def _split_ham_with_grouping(tape: qml.tape.QuantumScript): def _split_using_qwc_grouping( tape: qml.tape.QuantumScript, - single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]], - offsets: List[float], + single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]], + offsets: list[float], ): """Split tapes using group_observables in the Pauli module. @@ -449,8 +450,8 @@ def _split_using_qwc_grouping( def _split_using_wires_grouping( tape: qml.tape.QuantumScript, - single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]], - offsets: List[float], + single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]], + offsets: list[float], ): """Split tapes by grouping observables based on overlapping wires. @@ -579,8 +580,8 @@ def _split_all_multi_term_obs_mps(tape: qml.tape.QuantumScript): def _processing_fn_no_grouping( res: ResultBatch, - single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float]]], - offsets: List[float], + single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]], + offsets: list[float], shots: Shots, batch_size: int, ): @@ -631,9 +632,9 @@ def _processing_fn_no_grouping( def _processing_fn_with_grouping( res: ResultBatch, - single_term_obs_mps: Dict[MeasurementProcess, Tuple[List[int], List[float], int, int]], - offsets: List[float], - group_sizes: List[int], + single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float], int, int]], + offsets: list[float], + group_sizes: list[int], shots: Shots, batch_size: int, ): @@ -699,7 +700,7 @@ def _processing_fn_with_grouping( return tuple(res_for_each_mp) -def _sum_terms(res: ResultBatch, coeffs: List[float], offset: float, shape: Tuple) -> Result: +def _sum_terms(res: ResultBatch, coeffs: list[float], offset: float, shape: tuple) -> Result: """Sum results from measurements of multiple terms in a multi-term observable.""" # Trivially return the original result @@ -733,7 +734,7 @@ def _mp_to_obs(mp: MeasurementProcess, tape: qml.tape.QuantumScript) -> qml.oper return qml.prod(*(qml.Z(wire) for wire in obs_wires)) -def _infer_result_shape(shots: Shots, batch_size: int) -> Tuple: +def _infer_result_shape(shots: Shots, batch_size: int) -> tuple: """Based on the result, infer the ([,n_shots] [,batch_size]) shape of the result.""" shape = () diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py index b2d4cb38773..92837802e0f 100644 --- a/pennylane/transforms/transpile.py +++ b/pennylane/transforms/transpile.py @@ -3,7 +3,6 @@ """ from functools import partial -from typing import Callable, Sequence import networkx as nx @@ -13,9 +12,9 @@ from pennylane.ops import __all__ as all_ops from pennylane.ops.qubit import SWAP from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform -from pennylane.typing import Result, ResultBatch +from pennylane.typing import PostprocessingFn def state_transposition(results, mps, new_wire_order, original_wire_order): @@ -63,7 +62,7 @@ def _process_measurements(expanded_tape, device_wires, is_default_mixed): @transform def transpile( tape: QuantumTape, coupling_map, device=None -) -> tuple[Sequence[QuantumTape], Callable[[ResultBatch], Result]]: +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transpile a circuit according to a desired coupling map .. warning:: diff --git a/pennylane/transforms/unitary_to_rot.py b/pennylane/transforms/unitary_to_rot.py index 4b541e826d5..65d1341c740 100644 --- a/pennylane/transforms/unitary_to_rot.py +++ b/pennylane/transforms/unitary_to_rot.py @@ -14,17 +14,17 @@ """ A transform for decomposing arbitrary single-qubit QubitUnitary gates into elementary gates. """ -from typing import Callable, Sequence import pennylane as qml from pennylane.ops.op_math.decompositions import one_qubit_decomposition, two_qubit_decomposition from pennylane.queuing import QueuingManager -from pennylane.tape import QuantumTape +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform +from pennylane.typing import PostprocessingFn @transform -def unitary_to_rot(tape: QuantumTape) -> (Sequence[QuantumTape], Callable): +def unitary_to_rot(tape: QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: r"""Quantum function transform to decomposes all instances of single-qubit and select instances of two-qubit :class:`~.QubitUnitary` operations to parametrized single-qubit operations. diff --git a/pennylane/transforms/zx/converter.py b/pennylane/transforms/zx/converter.py index 0567b19bfac..af009a05940 100644 --- a/pennylane/transforms/zx/converter.py +++ b/pennylane/transforms/zx/converter.py @@ -16,14 +16,14 @@ from collections import OrderedDict from functools import partial -from typing import Callable, Sequence import numpy as np import pennylane as qml from pennylane.operation import Operator -from pennylane.tape import QuantumScript, QuantumTape +from pennylane.tape import QuantumScript, QuantumTape, QuantumTapeBatch from pennylane.transforms import TransformError, transform +from pennylane.typing import PostprocessingFn from pennylane.wires import Wires @@ -266,7 +266,7 @@ def mod_5_4(): @partial(transform, is_informative=True) def _to_zx_transform( tape: QuantumTape, expand_measurements=False -) -> (Sequence[QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Private function to convert a PennyLane tape to a `PyZX graph `_ .""" # Avoid to make PyZX a requirement for PennyLane. try: diff --git a/pennylane/typing.py b/pennylane/typing.py index 4a289d7c036..13f50d0b9cc 100644 --- a/pennylane/typing.py +++ b/pennylane/typing.py @@ -14,9 +14,10 @@ """This file contains different PennyLane types.""" import contextlib -# pylint: disable=import-outside-toplevel, too-few-public-methods +# pylint: disable=import-outside-toplevel, too-few-public-methods, unused-import import sys -from typing import Sequence, TypeVar, Union +from collections.abc import Callable, Sequence +from typing import TypeVar, Union import numpy as np from autograd.numpy.numpy_boxes import ArrayBox @@ -123,4 +124,7 @@ def _is_torch(other, subclass=False): ResultBatch = Sequence[Result] +PostprocessingFn = Callable[[ResultBatch], Result] +BatchPostprocessingFn = Callable[[ResultBatch], ResultBatch] + JSON = Union[None, int, str, bool, list["JSON"], dict[str, "JSON"]] diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py index 9c95c6003ac..5eda7a2d4e1 100644 --- a/pennylane/workflow/construct_batch.py +++ b/pennylane/workflow/construct_batch.py @@ -15,11 +15,14 @@ """ import inspect +from collections.abc import Callable from contextlib import nullcontext from functools import wraps -from typing import Callable, Literal, Optional, Tuple, Union +from typing import Literal, Optional, Union import pennylane as qml +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn from .qnode import QNode, _get_device_shots, _make_execution_config @@ -316,7 +319,7 @@ def circuit(x): """ # pylint: disable=protected-access - def batch_constructor(*args, **kwargs) -> Tuple[Tuple["qml.tape.QuantumTape", Callable]]: + def batch_constructor(*args, **kwargs) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Create a batch of tapes and a post processing function.""" if "shots" in inspect.signature(qnode.func).parameters: shots = _get_device_shots(qnode.device) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index b020cb7c7f3..cac6aafd2fe 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -24,15 +24,17 @@ import inspect import logging import warnings +from collections.abc import Callable, MutableMapping, Sequence from functools import partial -from typing import Callable, MutableMapping, Optional, Sequence, Tuple, Union +from typing import Optional, Union from cachetools import Cache, LRUCache import pennylane as qml -from pennylane.tape import QuantumTape +from pennylane.data.base.attribute import UNSET +from pennylane.tape import QuantumTape, QuantumTapeBatch from pennylane.transforms import transform -from pennylane.typing import ResultBatch +from pennylane.typing import PostprocessingFn, Result, ResultBatch from .jacobian_products import ( DeviceDerivatives, @@ -95,7 +97,7 @@ def _adjoint_jacobian_expansion( - tapes: Sequence[QuantumTape], grad_on_execution: bool, interface: str, max_expansion: int + tapes: QuantumTapeBatch, grad_on_execution: bool, interface: str, max_expansion: int ): """Performs adjoint jacobian specific expansion. Expands so that every trainable operation has a generator. @@ -180,12 +182,12 @@ def _get_ml_boundary_execute( def _batch_transform( - tapes: Sequence[QuantumTape], + tapes: QuantumTapeBatch, device: device_type, config: "qml.devices.ExecutionConfig", override_shots: Union[bool, int, Sequence[int]] = False, device_batch_transform: bool = True, -) -> Tuple[Sequence[QuantumTape], Callable, "qml.devices.ExecutionConfig"]: +) -> tuple[QuantumTapeBatch, PostprocessingFn, "qml.devices.ExecutionConfig"]: """Apply the device batch transform unless requested not to. Args: @@ -288,7 +290,7 @@ def _make_inner_execute( else: device_execution = partial(device.execute, execution_config=execution_config) - def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: + def inner_execute(tapes: QuantumTapeBatch, **_) -> ResultBatch: """Execution that occurs within a machine learning framework boundary. Closure Variables: @@ -331,7 +333,7 @@ def _cache_transform(tape: QuantumTape, cache: MutableMapping): This function makes use of :attr:`.QuantumTape.hash` to identify unique tapes. """ - def cache_hit_postprocessing(_results: Tuple[Tuple]) -> Tuple: + def cache_hit_postprocessing(_results: ResultBatch) -> Result: result = cache[tape.hash] if result is not None: if tape.shots and getattr(cache, "_persistent_cache", True): @@ -346,7 +348,7 @@ def cache_hit_postprocessing(_results: Tuple[Tuple]) -> Tuple: if tape.hash in cache: return [], cache_hit_postprocessing - def cache_miss_postprocessing(results: Tuple[Tuple]) -> Tuple: + def cache_miss_postprocessing(results: ResultBatch) -> Result: result = results[0] cache[tape.hash] = result return result @@ -408,8 +410,65 @@ def _get_interface_name(tapes, interface): return interface +def _deprecated_arguments_warnings( + tapes, override_shots, expand_fn, max_expansion, device_batch_transform +): + """Helper function to raise exceptions and pass codefactor checks regarding the length of the function""" + + if device_batch_transform is not None: + warnings.warn( + "The device_batch_transform argument is deprecated and will be removed in version 0.39. " + "Instead, please create a TransformProgram with the desired preprocessing and pass " + "it to the transform_program argument of qml.execute.", + qml.PennyLaneDeprecationWarning, + ) + else: + device_batch_transform = True + + if override_shots is not UNSET: + warnings.warn( + "The override_shots argument is deprecated and will be removed in version 0.39. " + "Instead, please add the shots to the QuantumTape's to be executed.", + qml.PennyLaneDeprecationWarning, + ) + if override_shots is not False: + tapes = tuple( + qml.tape.QuantumScript( + t.operations, + t.measurements, + trainable_params=t.trainable_params, + shots=override_shots, + ) + for t in tapes + ) + else: + override_shots = False + + if expand_fn is not UNSET: + warnings.warn( + "The expand_fn argument is deprecated and will be removed in version 0.39. " + "Instead, please create a TransformProgram with the desired preprocessing and pass " + "it to the transform_program argument of qml.execute.", + qml.PennyLaneDeprecationWarning, + ) + else: + expand_fn = "device" + + if max_expansion is not None: + warnings.warn( + "The max_expansion argument is deprecated and will be removed in version 0.39. " + "Instead, please use qml.devices.preprocess.decompose with the desired expansion level, " + "add it to a TransformProgram and pass it to the transform_program argument of qml.execute.", + qml.PennyLaneDeprecationWarning, + ) + else: + max_expansion = 10 + + return tapes, override_shots, expand_fn, max_expansion, device_batch_transform + + def execute( - tapes: Sequence[QuantumTape], + tapes: QuantumTapeBatch, device: device_type, gradient_fn: Optional[Union[Callable, str]] = None, interface="auto", @@ -421,10 +480,10 @@ def execute( cache: Union[None, bool, dict, Cache] = True, cachesize=10000, max_diff=1, - override_shots: int = False, - expand_fn="device", # type: ignore - max_expansion=10, - device_batch_transform=True, + override_shots: int = UNSET, + expand_fn=UNSET, # type: ignore + max_expansion=None, + device_batch_transform=None, device_vjp=False, mcm_config=None, ) -> ResultBatch: @@ -480,6 +539,53 @@ def execute( list[tensor_like[float]]: A nested list of tape results. Each element in the returned list corresponds in order to the provided tapes. + .. warning:: + + The following arguments are deprecated and will be removed in version 0.39: + ``expand_fn``, ``max_expansion``, and ``device_batch_transform``. + Instead, please create a :class:`~.TransformProgram` with the desired preprocessing and + pass it to the ``transform_program`` argument. For instance, we can create a program that uses + the ``qml.devices.preprocess.decompose`` transform with the desired expansion level and pass it + to the ``qml.execute`` function: + + .. code-block:: python + + from pennylane.devices.preprocess import decompose + from pennylane.transforms.core import TransformProgram + + def stopping_condition(obj): + return obj.name in {"CNOT", "RX", "RZ"} + + tape = qml.tape.QuantumScript([qml.IsingXX(1.2, wires=(0,1))], [qml.expval(qml.Z(0))]) + + program = TransformProgram() + program.add_transform( + decompose, + stopping_condition=stopping_condition, + max_expansion=10, + ) + + dev = qml.device("default.qubit", wires=2) + + >>> qml.execute([tape], dev, transform_program=program) + (0.36235775447667357,) + + .. warning:: + + The ``override_shots`` argument is deprecated and will be removed in version 0.39. + Instead, please add the shots to the ``QuantumTape``'s to be executed. For instance: + + .. code-block:: python + + dev = qml.device("default.qubit", wires=1) + operations = [qml.PauliX(0)] + measurements = [qml.expval(qml.PauliZ(0))] + qs = qml.tape.QuantumTape(operations, measurements, shots=100) + + >>> qml.execute([qs], dev) + (-1.0,) + + **Example** Consider the following cost function: @@ -562,6 +668,12 @@ def cost_fn(params, x): "::L".join(str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]), ) + tapes, override_shots, expand_fn, max_expansion, device_batch_transform = ( + _deprecated_arguments_warnings( + tapes, override_shots, expand_fn, max_expansion, device_batch_transform + ) + ) + ### Specifying and preprocessing variables #### interface = _get_interface_name(tapes, interface) @@ -710,7 +822,7 @@ def execute_fn(internal_tapes): else: - def execute_fn(internal_tapes) -> Tuple[ResultBatch, Tuple]: + def execute_fn(internal_tapes) -> tuple[ResultBatch, tuple]: """A wrapper around device.execute that adds an empty tuple instead of derivatives. Closure Variables: diff --git a/pennylane/workflow/interfaces/autograd.py b/pennylane/workflow/interfaces/autograd.py index f3d8eda83b8..2b2e0ff8f4b 100644 --- a/pennylane/workflow/interfaces/autograd.py +++ b/pennylane/workflow/interfaces/autograd.py @@ -83,15 +83,15 @@ def grad_fn(dy): """ # pylint: disable=too-many-arguments, unused-argument import logging -from typing import Callable, Tuple +from collections.abc import Callable import autograd from autograd.numpy.numpy_boxes import ArrayBox import pennylane as qml +from pennylane.tape import QuantumTapeBatch -Batch = Tuple[qml.tape.QuantumTape] -ExecuteFn = Callable[[Batch], qml.typing.ResultBatch] +ExecuteFn = Callable[[QuantumTapeBatch], qml.typing.ResultBatch] logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -99,7 +99,7 @@ def grad_fn(dy): # pylint: disable=unused-argument def autograd_execute( - tapes: Batch, + tapes: QuantumTapeBatch, execute_fn: ExecuteFn, jpc: qml.workflow.jacobian_products.JacobianProductCalculator, device=None, diff --git a/pennylane/workflow/interfaces/jax.py b/pennylane/workflow/interfaces/jax.py index 141a4f99dab..5fe6282b10e 100644 --- a/pennylane/workflow/interfaces/jax.py +++ b/pennylane/workflow/interfaces/jax.py @@ -127,12 +127,13 @@ def f_and_jvp(primals, tangents): # pylint: disable=unused-argument import logging -from typing import Callable, Tuple +from collections.abc import Callable import jax import jax.numpy as jnp import pennylane as qml +from pennylane.tape import QuantumTapeBatch from pennylane.transforms import convert_to_numpy_parameters from pennylane.typing import ResultBatch @@ -142,8 +143,7 @@ def f_and_jvp(primals, tangents): logger.addHandler(logging.NullHandler()) -Batch = Tuple[qml.tape.QuantumTape] -ExecuteFn = Callable[[Batch], qml.typing.ResultBatch] +ExecuteFn = Callable[[QuantumTapeBatch], qml.typing.ResultBatch] @dataclasses.dataclass @@ -161,7 +161,7 @@ class _NonPytreeWrapper: """ - vals: Batch = None + vals: QuantumTapeBatch = None def _set_copy_and_unwrap_tape(t, a, unwrap=True): @@ -243,7 +243,7 @@ def _execute_and_compute_jvp(tapes, execute_fn, jpc, primals, tangents): _execute_jvp.defjvp(_execute_and_compute_jvp) -def jax_jvp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc, device=None): +def jax_jvp_execute(tapes: QuantumTapeBatch, execute_fn: ExecuteFn, jpc, device=None): """Execute a batch of tapes with JAX parameters using JVP derivatives. Args: diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index e8d9832e61c..8a93ce1de00 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -17,17 +17,16 @@ import abc import inspect import logging -from typing import Callable, Optional, Tuple, Union +from collections.abc import Callable, Sequence +from typing import Optional, Union import numpy as np from cachetools import LRUCache import pennylane as qml -from pennylane.tape import QuantumScript +from pennylane.tape import QuantumTapeBatch from pennylane.typing import ResultBatch, TensorLike -Batch = Tuple[QuantumScript] - logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -77,14 +76,14 @@ class JacobianProductCalculator(abc.ABC): @abc.abstractmethod def execute_and_compute_jvp( - self, tapes: Batch, tangents: Tuple[Tuple[TensorLike]] - ) -> Tuple[ResultBatch, Tuple]: + self, tapes: QuantumTapeBatch, tangents: Sequence[Sequence[TensorLike]] + ) -> tuple[ResultBatch, tuple]: """Calculate both the results for a batch of tapes and the jvp. This method is required to compute JVPs in the JAX interface. Args: - tapes (tuple[.QuantumScript]): The batch of tapes to take the derivatives of + tapes (Sequence[.QuantumScript | .QuantumTape]): The batch of tapes to take the derivatives of tangents (Sequence[Sequence[TensorLike]]): the tangents for the parameters of the tape. The ``i`` th tangent corresponds to the ``i`` th tape, and the ``j`` th entry into a tangent entry corresponds to the ``j`` th trainable parameter of the tape. @@ -118,7 +117,7 @@ def execute_and_compute_jvp( """ @abc.abstractmethod - def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]) -> Tuple: + def compute_vjp(self, tapes: QuantumTapeBatch, dy: Sequence[Sequence[TensorLike]]) -> tuple: """Compute the vjp for a given batch of tapes. This method is used by autograd, torch, and tensorflow to compute VJPs. @@ -158,7 +157,7 @@ def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]) -> Tuple: """ @abc.abstractmethod - def compute_jacobian(self, tapes: Batch) -> Tuple: + def compute_jacobian(self, tapes: QuantumTapeBatch) -> tuple: """Compute the full Jacobian for a batch of tapes. This method is required to compute Jacobians in the ``tensorflow`` interface @@ -182,7 +181,7 @@ def compute_jacobian(self, tapes: Batch) -> Tuple: """ @abc.abstractmethod - def execute_and_compute_jacobian(self, tapes: Batch) -> Tuple: + def execute_and_compute_jacobian(self, tapes: QuantumTapeBatch) -> tuple[ResultBatch, tuple]: """Compute the results and the full Jacobian for a batch of tapes. This method is required to compute Jacobians in the ``jax-jit`` interface @@ -241,7 +240,7 @@ def __repr__(self): def __init__( self, inner_execute: Callable, - gradient_transform: "pennylane.transforms.core.TransformDispatcher", + gradient_transform: "qml.transforms.core.TransformDispatcher", gradient_kwargs: Optional[dict] = None, cache_full_jacobian: bool = False, ): @@ -265,7 +264,9 @@ def __init__( self._cache_full_jacobian = cache_full_jacobian self._cache = LRUCache(maxsize=10) - def execute_and_compute_jvp(self, tapes: Batch, tangents: Tuple[Tuple[TensorLike]]): + def execute_and_compute_jvp( + self, tapes: QuantumTapeBatch, tangents: Sequence[Sequence[TensorLike]] + ): if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("execute_and_compute_jvp called with (%s, %s)", tapes, tangents) @@ -288,7 +289,7 @@ def execute_and_compute_jvp(self, tapes: Batch, tangents: Tuple[Tuple[TensorLike jvps = jvp_processing_fn(jvp_results) return tuple(results), tuple(jvps) - def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]): + def compute_vjp(self, tapes: QuantumTapeBatch, dy: Sequence[Sequence[TensorLike]]): if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("compute_vjp called with (%s, %s)", tapes, dy) @@ -303,7 +304,7 @@ def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]): vjp_results = self._inner_execute(tuple(vjp_tapes)) return tuple(processing_fn(vjp_results)) - def execute_and_compute_jacobian(self, tapes: Batch): + def execute_and_compute_jacobian(self, tapes: QuantumTapeBatch): if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("execute_and_compute_jacobian called with %s", tapes) @@ -318,7 +319,7 @@ def execute_and_compute_jacobian(self, tapes: Batch): jacs = jac_postprocessing(jac_results) return tuple(results), tuple(jacs) - def compute_jacobian(self, tapes: Batch): + def compute_jacobian(self, tapes: QuantumTapeBatch): if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("compute_jacobian called with %s", tapes) if tapes in self._cache: @@ -420,7 +421,7 @@ def __init__( self._results_cache = LRUCache(maxsize=10) self._jacs_cache = LRUCache(maxsize=10) - def _dev_execute_and_compute_derivatives(self, tapes: Batch): + def _dev_execute_and_compute_derivatives(self, tapes: QuantumTapeBatch): """ Converts tapes to numpy before computing the the results and derivatives on the device. @@ -431,7 +432,7 @@ def _dev_execute_and_compute_derivatives(self, tapes: Batch): return self._device.execute_and_compute_derivatives(numpy_tapes, self._execution_config) return self._device.execute_and_gradients(numpy_tapes, **self._gradient_kwargs) - def _dev_execute(self, tapes: Batch): + def _dev_execute(self, tapes: QuantumTapeBatch): """ Converts tapes to numpy before computing just the results on the device. @@ -442,7 +443,7 @@ def _dev_execute(self, tapes: Batch): return self._device.execute(numpy_tapes, self._execution_config) return self._device.batch_execute(numpy_tapes) - def _dev_compute_derivatives(self, tapes: Batch): + def _dev_compute_derivatives(self, tapes: QuantumTapeBatch): """ Converts tapes to numpy before computing the derivatives on the device. @@ -453,7 +454,7 @@ def _dev_compute_derivatives(self, tapes: Batch): return self._device.compute_derivatives(numpy_tapes, self._execution_config) return self._device.gradients(numpy_tapes, **self._gradient_kwargs) - def execute_and_cache_jacobian(self, tapes: Batch): + def execute_and_cache_jacobian(self, tapes: QuantumTapeBatch): """Forward pass used to cache the results and jacobians. Args: @@ -473,7 +474,7 @@ def execute_and_cache_jacobian(self, tapes: Batch): self._jacs_cache[tapes] = jac return results - def execute_and_compute_jvp(self, tapes: Batch, tangents): + def execute_and_compute_jvp(self, tapes: QuantumTapeBatch, tangents): """Calculate both the results for a batch of tapes and the jvp. This method is required to compute JVPs in the JAX interface. @@ -665,15 +666,15 @@ def __init__( self._execution_config = execution_config def execute_and_compute_jvp( - self, tapes: Batch, tangents: Tuple[Tuple[TensorLike]] - ) -> Tuple[ResultBatch, Tuple]: + self, tapes: QuantumTapeBatch, tangents: Sequence[Sequence[TensorLike]] + ) -> tuple[ResultBatch, tuple]: if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("execute_and_compute_jvp called with (%s, %s)", tapes, tangents) numpy_tapes, _ = qml.transforms.convert_to_numpy_parameters(tapes) tangents = qml.math.unwrap(tangents) return self._device.execute_and_compute_jvp(numpy_tapes, tangents, self._execution_config) - def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]) -> Tuple: + def compute_vjp(self, tapes: QuantumTapeBatch, dy: Sequence[Sequence[TensorLike]]) -> tuple: if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("compute_vjp called with (%s, %s)", tapes, dy) numpy_tapes, _ = qml.transforms.convert_to_numpy_parameters(tapes) @@ -687,13 +688,13 @@ def compute_vjp(self, tapes: Batch, dy: Tuple[Tuple[TensorLike]]) -> Tuple: res.append(r) return res - def compute_jacobian(self, tapes: Batch): + def compute_jacobian(self, tapes: QuantumTapeBatch): if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("compute_jacobian called with %s", tapes) numpy_tapes, _ = qml.transforms.convert_to_numpy_parameters(tapes) return self._device.compute_derivatives(numpy_tapes, self._execution_config) - def execute_and_compute_jacobian(self, tapes: Batch) -> Tuple: + def execute_and_compute_jacobian(self, tapes: QuantumTapeBatch) -> tuple: if logger.isEnabledFor(logging.DEBUG): # pragma: no cover logger.debug("execute_and_compute_jacobian called with %s", tapes) numpy_tapes, _ = qml.transforms.convert_to_numpy_parameters(tapes) diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py index fe807e403e8..09c667bdd3e 100644 --- a/pennylane/workflow/qnode.py +++ b/pennylane/workflow/qnode.py @@ -106,7 +106,7 @@ def _to_qfunc_output_type( class QNode: - """Represents a quantum node in the hybrid computational graph. + r"""Represents a quantum node in the hybrid computational graph. A *quantum node* contains a :ref:`quantum function ` (corresponding to a `variational circuit `) @@ -238,6 +238,14 @@ class QNode: method. Please refer to the :mod:`qml.gradients <.gradients>` module for details on supported options for your chosen gradient transform. + .. warning:: + + The ``expansion_strategy`` argument is deprecated and will be removed in version 0.39. + + .. warning:: + + The ``max_expansion`` argument is deprecated and will be removed in version 0.39. + **Example** QNodes can be created by decorating a quantum function: @@ -448,8 +456,8 @@ def __init__( device: Union[Device, "qml.devices.Device"], interface="auto", diff_method="best", - expansion_strategy="gradient", - max_expansion=10, + expansion_strategy=None, + max_expansion=None, grad_on_execution="best", cache="auto", cachesize=10000, @@ -459,6 +467,15 @@ def __init__( mcm_method=None, **gradient_kwargs, ): + # Moving it here since the old default value is checked on debugging + if max_expansion is not None: + warnings.warn( + "The max_expansion argument is deprecated and will be removed in version 0.39. ", + qml.PennyLaneDeprecationWarning, + ) + else: + max_expansion = 10 + if logger.isEnabledFor(logging.DEBUG): logger.debug( """Creating QNode(func=%s, device=%s, interface=%s, diff_method=%s, expansion_strategy=%s, max_expansion=%s, grad_on_execution=%s, cache=%s, cachesize=%s, max_diff=%s, gradient_kwargs=%s""", @@ -479,6 +496,17 @@ def __init__( gradient_kwargs, ) + if expansion_strategy is not None: + warnings.warn( + "The 'expansion_strategy' attribute is deprecated and will be removed " + "in version 0.39. For full control over the stage to which the tape is " + "constructed, use the 'pennylane.workflow.construct_batch' function.", + qml.PennyLaneDeprecationWarning, + ) + + # Default to "gradient" to maintain default behaviour of "draw" and "specs" + expansion_strategy = expansion_strategy or "gradient" + if interface not in SUPPORTED_INTERFACES: raise qml.QuantumFunctionError( f"Unknown interface {interface}. Interface must be " @@ -1081,7 +1109,6 @@ def _execution_component(self, args: tuple, kwargs: dict, override_shots) -> qml mcm_config=mcm_config, interface=self.interface, ) - override_shots = 1 elif hasattr(self.device, "capabilities"): inner_transform_program.add_transform( qml.defer_measurements, @@ -1099,19 +1126,26 @@ def _execution_component(self, args: tuple, kwargs: dict, override_shots) -> qml full_transform_program.set_classical_component(self, args, kwargs) _prune_dynamic_transform(full_transform_program, inner_transform_program) - # pylint: disable=unexpected-keyword-arg - res = qml.execute( - (self._tape,), - device=self.device, - gradient_fn=self.gradient_fn, - interface=self.interface, - transform_program=full_transform_program, - inner_transform=inner_transform_program, - config=config, - gradient_kwargs=self.gradient_kwargs, - override_shots=override_shots, - **self.execute_kwargs, - ) + with warnings.catch_warnings(): + # TODO: remove this once the cycle for the arguements have finished, i.e. 0.39. + warnings.filterwarnings( + action="ignore", + message=r".*argument is deprecated and will be removed in version 0.39.*", + category=qml.PennyLaneDeprecationWarning, + ) + # pylint: disable=unexpected-keyword-arg + res = qml.execute( + (self._tape,), + device=self.device, + gradient_fn=self.gradient_fn, + interface=self.interface, + transform_program=full_transform_program, + inner_transform=inner_transform_program, + config=config, + gradient_kwargs=self.gradient_kwargs, + override_shots=override_shots, + **self.execute_kwargs, + ) res = res[0] # convert result to the interface in case the qfunc has no parameters diff --git a/pennylane/workflow/set_shots.py b/pennylane/workflow/set_shots.py index bb463233fed..1bc7dff7f33 100644 --- a/pennylane/workflow/set_shots.py +++ b/pennylane/workflow/set_shots.py @@ -24,15 +24,14 @@ @contextlib.contextmanager def set_shots(device, shots): - """Context manager to temporarily change the shots - of a device. + r"""Context manager to temporarily change the shots of a device. This context manager can be used in two ways. As a standard context manager: >>> dev = qml.device("default.qubit.legacy", wires=2, shots=None) - >>> with set_shots(dev, shots=100): + >>> with qml.workflow.set_shots(dev, shots=100): ... print(dev.shots) 100 >>> print(dev.shots) @@ -40,7 +39,7 @@ def set_shots(device, shots): Or as a decorator that acts on a function that uses the device: - >>> set_shots(dev, shots=100)(lambda: dev.shots)() + >>> qml.workflow.set_shots(dev, shots=100)(lambda: dev.shots)() 100 """ if isinstance(device, qml.devices.Device): diff --git a/tests/circuit_graph/test_circuit_graph.py b/tests/circuit_graph/test_circuit_graph.py index df367b7af63..10886f559f2 100644 --- a/tests/circuit_graph/test_circuit_graph.py +++ b/tests/circuit_graph/test_circuit_graph.py @@ -139,31 +139,19 @@ def test_dependence(self, ops, obs): assert len(graph.node_indexes()) == 9 assert len(graph.edges()) == 9 - queue = ops + obs - - # all ops should be nodes in the graph - for k in queue: - assert k in graph.nodes() - - # all nodes in the graph should be ops - # for k in graph.nodes: - for k in graph.nodes(): - assert k is queue[k.queue_idx] - a = set((graph.get_node_data(e[0]), graph.get_node_data(e[1])) for e in graph.edge_list()) - b = set( - (queue[a], queue[b]) - for a, b in [ - (0, 3), - (1, 3), - (2, 4), - (3, 5), - (3, 6), - (4, 5), - (5, 7), - (5, 8), - (6, 8), - ] - ) + a = {(graph.get_node_data(e[0]), graph.get_node_data(e[1])) for e in graph.edge_list()} + + b = { + (0, 3), + (1, 3), + (2, 4), + (3, 5), + (3, 6), + (4, 5), + (5, 7), + (5, 8), + (6, 8), + } assert a == b def test_ancestors_and_descendants_example(self, ops, obs): @@ -182,25 +170,44 @@ def test_ancestors_and_descendants_example(self, ops, obs): descendants = circuit.descendants([queue[6]]) assert descendants == [queue[8]] - def test_in_topological_order_example(self, ops, obs): - """ - Test ``_in_topological_order`` method returns the expected result. - """ - circuit = CircuitGraph(ops, obs, Wires([0, 1, 2])) - - to = circuit._in_topological_order(ops) - - to_expected = [ - qml.RZ(0.35, wires=[2]), - qml.Hadamard(wires=[2]), - qml.RY(0.35, wires=[1]), - qml.RX(0.43, wires=[0]), - qml.CNOT(wires=[0, 1]), - qml.PauliX(wires=[1]), - qml.CNOT(wires=[2, 0]), - ] - - assert str(to) == str(to_expected) + def test_ancestors_and_descendents_repeated_op(self): + """Test ancestors and descendents raises a ValueError is the requested operation occurs more than once.""" + + op = qml.X(0) + ops = [op, qml.Y(0), op, qml.Z(0), op] + graph = CircuitGraph(ops, [], [0, 1, 2]) + + with pytest.raises(ValueError, match=r"operator that occurs multiple times."): + graph.ancestors([op]) + with pytest.raises(ValueError, match=r"operator that occurs multiple times."): + graph.descendants([op]) + with pytest.raises(ValueError, match=r"operator that occurs multiple times."): + graph.ancestors_in_order([op]) + with pytest.raises(ValueError, match=r"operator that occurs multiple times."): + graph.descendants_in_order([op]) + + def test_ancestors_and_descendents_single_op_error(self): + """Test ancestors and descendents raises a ValueError is the requested operation occurs more than once.""" + + op = qml.Z(0) + graph = CircuitGraph([op], [], [0, 1, 2]) + + with pytest.raises( + ValueError, match=r"CircuitGraph.ancestors accepts an iterable of operators" + ): + graph.ancestors(op) + with pytest.raises( + ValueError, match=r"CircuitGraph.descendants accepts an iterable of operators" + ): + graph.descendants(op) + with pytest.raises( + ValueError, match=r"CircuitGraph.ancestors accepts an iterable of operators" + ): + graph.ancestors_in_order(op) + with pytest.raises( + ValueError, match=r"CircuitGraph.descendants accepts an iterable of operators" + ): + graph.descendants_in_order(op) def test_update_node(self, ops, obs): """Changing nodes in the graph.""" @@ -209,6 +216,9 @@ def test_update_node(self, ops, obs): new = qml.RX(0.1, wires=0) circuit.update_node(ops[0], new) assert circuit.operations[0] is new + new_mp = qml.var(qml.Y(0)) + circuit.update_node(obs[0], new_mp) + assert circuit.observables[0] is new_mp def test_update_node_error(self, ops, obs): """Test that changing nodes in the graph may raise an error.""" @@ -300,43 +310,6 @@ def test_max_simultaneous_measurements(self, circ, expected): circuit = qnode.qtape.graph assert circuit.max_simultaneous_measurements == expected - def test_grid_when_sample_no_wires(self): - """A test to ensure the sample operation applies to all wires on the - `CircuitGraph`s grid when none are explicitly provided.""" - - ops = [qml.Hadamard(wires=0), qml.CNOT(wires=[0, 1])] - obs_no_wires = [qml.sample(op=None, wires=None)] - obs_w_wires = [qml.sample(op=None, wires=[0, 1, 2])] - - circuit_no_wires = CircuitGraph(ops, obs_no_wires, wires=Wires([0, 1, 2])) - circuit_w_wires = CircuitGraph(ops, obs_w_wires, wires=Wires([0, 1, 2])) - - sample_w_wires_op = qml.sample(op=None, wires=[0, 1, 2]) - expected_grid_w_wires = { - 0: [ops[0], ops[1], sample_w_wires_op], - 1: [ops[1], sample_w_wires_op], - 2: [sample_w_wires_op], - } - - sample_no_wires_op = qml.sample(op=None, wires=None) - expected_grid_no_wires = { - 0: [ops[0], ops[1], sample_no_wires_op], - 1: [ops[1], sample_no_wires_op], - 2: [sample_no_wires_op], - } - - for key in range(3): - lst_w_wires = circuit_w_wires._grid[key] - lst_no_wires = circuit_no_wires._grid[key] - lst_expected_w_wires = expected_grid_w_wires[key] - lst_expected_no_wires = expected_grid_no_wires[key] - - for el1, el2 in zip(lst_w_wires, lst_expected_w_wires): - qml.assert_equal(el1, el2) - - for el1, el2 in zip(lst_no_wires, lst_expected_no_wires): - qml.assert_equal(el1, el2) - def test_print_contents(self): """Tests if the circuit prints correct.""" ops = [qml.Hadamard(wires=0), qml.CNOT(wires=[0, 1])] @@ -354,6 +327,7 @@ def test_print_contents(self): tape_depth = ( ([qml.PauliZ(0), qml.CNOT([0, 1]), qml.RX(1.23, 2)], 2), + ([qml.X(0)] * 4, 4), ([qml.Hadamard(0), qml.CNOT([0, 1]), CustomOpDepth3(wires=[1, 0])], 5), ( [ @@ -381,3 +355,34 @@ def test_get_depth(self, ops, true_depth): """Test that depth is computed correctly for operations that define a custom depth > 1""" cg = CircuitGraph(ops, [], wires=[0, 1, 2, 3, 4]) assert cg.get_depth() == true_depth + + +def test_has_path(): + """Test has_path and has_path_idx.""" + + ops = [qml.X(0), qml.X(3), qml.CNOT((0, 1)), qml.X(1), qml.X(3)] + graph = CircuitGraph(ops, [], wires=[0, 1, 2, 3, 4, 5]) + + assert graph.has_path(ops[0], ops[2]) + assert graph.has_path_idx(0, 2) + assert not graph.has_path(ops[0], ops[4]) + assert not graph.has_path_idx(0, 4) + + +def test_has_path_repeated_ops(): + """Test has_path and has_path_idx when an operation is repeated.""" + + op = qml.X(0) + ops = [op, qml.CNOT((0, 1)), op, qml.Y(1)] + + graph = CircuitGraph(ops, [], [0, 1, 2, 3]) + + assert graph.has_path_idx(0, 3) + assert graph.has_path_idx(1, 2) + with pytest.raises(ValueError, match="does not work with operations that have been repeated. "): + graph.has_path(op, ops[3]) + with pytest.raises(ValueError, match="does not work with operations that have been repeated. "): + graph.has_path(ops[1], op) + + # still works if they are the same operation. + assert graph.has_path(op, op) diff --git a/tests/conftest.py b/tests/conftest.py index 4c8eba9dc6d..3094a2fd784 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,7 @@ import os import pathlib import sys +import warnings import numpy as np import pytest @@ -49,6 +50,23 @@ def set_numpy_seed(): yield +@pytest.fixture(scope="function", autouse=True) +def capture_legacy_device_deprecation_warnings(): + with warnings.catch_warnings(record=True) as recwarn: + warnings.simplefilter("always") + yield + + for w in recwarn: + if isinstance(w, qml.PennyLaneDeprecationWarning): + assert "Use of 'default.qubit." in str(w.message) + assert "is deprecated" in str(w.message) + assert "use 'default.qubit'" in str(w.message) + + for w in recwarn: + if "Use of 'default.qubit." not in str(w.message): + warnings.warn(message=w.message, category=w.category) + + @pytest.fixture(scope="session") def tol(): """Numerical tolerance for equality tests.""" @@ -81,7 +99,8 @@ def n_subsystems_fixture(request): @pytest.fixture(scope="session") def qubit_device(n_subsystems): - return qml.device("default.qubit.legacy", wires=n_subsystems) + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Use of 'default.qubit.legacy'"): + return qml.device("default.qubit.legacy", wires=n_subsystems) @pytest.fixture(scope="function", params=[(np.float32, np.complex64), (np.float64, np.complex128)]) diff --git a/tests/data/base/test_attribute.py b/tests/data/base/test_attribute.py index 7cc0ee735c5..d38249c1672 100644 --- a/tests/data/base/test_attribute.py +++ b/tests/data/base/test_attribute.py @@ -17,8 +17,9 @@ # pylint: disable=unused-argument,unused-variable,too-many-public-methods +from collections.abc import Iterable from copy import copy, deepcopy -from typing import Any, Iterable, List +from typing import Any import numpy as np import pytest @@ -47,7 +48,7 @@ pytest.importorskip("h5py") -def _sort_types(types: Iterable[type]) -> List[type]: +def _sort_types(types: Iterable[type]) -> list[type]: """ pytest-split requires that test parameters are always in the same order between runs. This function ensures that collections of types @@ -126,7 +127,7 @@ class TestAttributeInfo: [ ({"py_type": None}, "None"), ({"py_type": "None"}, "None"), - ({"py_type": List[None]}, "list[None]"), + ({"py_type": list[None]}, "list[None]"), ({}, None), ], ) diff --git a/tests/data/base/test_typing_util.py b/tests/data/base/test_typing_util.py index bd9affb9cfa..aad94132c79 100644 --- a/tests/data/base/test_typing_util.py +++ b/tests/data/base/test_typing_util.py @@ -15,7 +15,7 @@ Tests for the :mod:`pennylane.data.base.typing_util` functions. """ -import typing +from typing import Optional, Type, Union import pytest @@ -30,16 +30,16 @@ "type_, expect", [ (list, "list"), - (typing.List, "list"), + (list, "list"), (Molecule, "pennylane.qchem.molecule.Molecule"), ("nonsense", "nonsense"), - (typing.List[int], "list[int]"), - (typing.List[typing.Tuple[int, "str"]], "list[tuple[int, str]]"), - (typing.Optional[int], "Union[int, None]"), - (typing.Union[int, "str", Molecule], "Union[int, str, pennylane.qchem.molecule.Molecule]"), + (list[int], "list[int]"), + (list[tuple[int, "str"]], "list[tuple[int, str]]"), + (Optional[int], "Union[int, None]"), + (Union[int, "str", Molecule], "Union[int, str, pennylane.qchem.molecule.Molecule]"), (str, "str"), - (typing.Type[str], "type[str]"), - (typing.Union[typing.List[typing.List[int]], str], "Union[list[list[int]], str]"), + (Type[str], "type[str]"), + (Union[list[list[int]], str], "Union[list[list[int]], str]"), ], ) def test_get_type_str(type_, expect): @@ -53,8 +53,8 @@ def test_get_type_str(type_, expect): [ (list, list), ([1, 2], list), - (typing.List, list), - (typing.List[int], list), + (list, list), + (list[int], list), (qml.RX, qml.RX), (qml.RX(1, [1]), qml.RX), ], @@ -70,7 +70,7 @@ def test_unset_bool(): assert not UNSET -@pytest.mark.parametrize("type_, expect", [(typing.List[typing.List[int]], (list, [list]))]) +@pytest.mark.parametrize("type_, expect", [(list[list[int]], (list, [list]))]) def test_resolve_special_type(type_, expect): """Test resolve_special_type().""" assert resolve_special_type(type_) == expect diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index bb0f7a11584..69e8fecdfd9 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -19,7 +19,7 @@ from numbers import Number from pathlib import Path -from typing import ClassVar, Dict +from typing import ClassVar import numpy as np import pytest @@ -406,7 +406,7 @@ class NewDataset( x: int y: str = field(doc="y documentation") - jsonable: Dict[str, str] = field(DatasetJSON, doc="json data") + jsonable: dict[str, str] = field(DatasetJSON, doc="json data") ds = NewDataset(init_arg=3, x=1, y="abc", jsonable={"a": "b"}) diff --git a/tests/devices/default_qubit/test_default_qubit_native_mcm.py b/tests/devices/default_qubit/test_default_qubit_native_mcm.py index a79ef70fbaa..be0f8574d97 100644 --- a/tests/devices/default_qubit/test_default_qubit_native_mcm.py +++ b/tests/devices/default_qubit/test_default_qubit_native_mcm.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Tests for default qubit preprocessing.""" -from typing import Sequence +from collections.abc import Sequence import mcm_utils import numpy as np diff --git a/tests/devices/default_tensor/test_default_tensor.py b/tests/devices/default_tensor/test_default_tensor.py index 9fb622e3174..203680053d4 100644 --- a/tests/devices/default_tensor/test_default_tensor.py +++ b/tests/devices/default_tensor/test_default_tensor.py @@ -405,7 +405,7 @@ def test_jax(self, method): jax = pytest.importorskip("jax") dev = qml.device("default.tensor", wires=1, method=method) - ref_dev = qml.device("default.qubit.jax", wires=1) + ref_dev = qml.device("default.qubit", wires=1) def circuit(x): qml.RX(x[1], wires=0) diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py index ec9fd6e65f5..be8e586adcc 100644 --- a/tests/devices/test_default_qubit_legacy.py +++ b/tests/devices/test_default_qubit_legacy.py @@ -2001,7 +2001,8 @@ class TestApplyOps: gates in DefaultQubitLegacy.""" state = np.arange(2**4, dtype=np.complex128).reshape((2, 2, 2, 2)) - dev = qml.device("default.qubit.legacy", wires=4) + with pytest.warns(qml.PennyLaneDeprecationWarning): + dev = qml.device("default.qubit.legacy", wires=4) single_qubit_ops = [ (qml.PauliX, dev._apply_x), diff --git a/tests/devices/test_default_qubit_legacy_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py index ebf21502971..14d741d0423 100644 --- a/tests/devices/test_default_qubit_legacy_broadcasting.py +++ b/tests/devices/test_default_qubit_legacy_broadcasting.py @@ -1619,7 +1619,8 @@ class TestApplyOpsBroadcasted: gates in DefaultQubitLegacy.""" broadcasted_state = np.arange(2**4 * 3, dtype=np.complex128).reshape((3, 2, 2, 2, 2)) - dev = qml.device("default.qubit.legacy", wires=4) + with pytest.warns(qml.PennyLaneDeprecationWarning): + dev = qml.device("default.qubit.legacy", wires=4) single_qubit_ops = [ (qml.PauliX, dev._apply_x), diff --git a/tests/test_device.py b/tests/devices/test_device.py similarity index 97% rename from tests/test_device.py rename to tests/devices/test_device.py index 1639f76ed57..a5f6631eb90 100644 --- a/tests/test_device.py +++ b/tests/devices/test_device.py @@ -188,6 +188,12 @@ def get_device(wires=1): yield get_device +# pylint: disable=pointless-statement +def test_invalid_attribute_in_devices_raises_error(): + with pytest.raises(AttributeError, match="'pennylane.devices' has no attribute 'blabla'"): + qml.devices.blabla + + def test_gradients_record(): """Test that execute_and_gradients and gradient both track the number of gradients requested.""" @@ -642,7 +648,7 @@ def test_default_expand_with_initial_state(self, op, decomp): prep = [op] ops = [qml.AngleEmbedding(features=[0.1], wires=[0], rotation="Z"), op, qml.PauliZ(wires=2)] - dev = qml.device("default.qubit.legacy", wires=3) + dev = qml.device("default.mixed", wires=3) tape = qml.tape.QuantumTape(ops=prep + ops, measurements=[], shots=100) new_tape = dev.default_expand_fn(tape) @@ -998,7 +1004,14 @@ def test_outdated_API(self, monkeypatch): with monkeypatch.context() as m: m.setattr(qml, "version", lambda: "0.0.1") with pytest.raises(DeviceError, match="plugin requires PennyLane versions"): - qml.device("default.qubit.legacy", wires=0) + qml.device("default.mixed", wires=0) + + def test_plugin_devices_from_devices_triggers_getattr(self, mocker): + spied = mocker.spy(qml.devices, "__getattr__") + + qml.devices.plugin_devices + + spied.assert_called_once() def test_refresh_entrypoints(self, monkeypatch): """Test that new entrypoints are found by the refresh_devices function""" @@ -1011,6 +1024,7 @@ def test_refresh_entrypoints(self, monkeypatch): # reimporting PennyLane within the context sets qml.plugin_devices to {} reload(qml) + reload(qml.devices.device_constructor) # since there are no entry points, there will be no plugin devices assert not qml.plugin_devices @@ -1023,6 +1037,7 @@ def test_refresh_entrypoints(self, monkeypatch): # Test teardown: re-import PennyLane to revert all changes and # restore the plugin_device dictionary reload(qml) + reload(qml.devices.device_constructor) def test_hot_refresh_entrypoints(self, monkeypatch): """Test that new entrypoints are found by the device loader if not currently present""" @@ -1034,9 +1049,10 @@ def test_hot_refresh_entrypoints(self, monkeypatch): m.setattr(metadata, "entry_points", lambda **kwargs: retval) # reimporting PennyLane within the context sets qml.plugin_devices to {} - reload(qml) + reload(qml.devices) + reload(qml.devices.device_constructor) - m.setattr(qml, "refresh_devices", lambda: None) + m.setattr(qml.devices.device_constructor, "refresh_devices", lambda: None) assert not qml.plugin_devices # since there are no entry points, there will be no plugin devices @@ -1052,10 +1068,11 @@ def test_hot_refresh_entrypoints(self, monkeypatch): # Test teardown: re-import PennyLane to revert all changes and # restore the plugin_device dictionary reload(qml) + reload(qml.devices.device_constructor) def test_shot_vector_property(self): """Tests shot vector initialization.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=[1, 3, 3, 4, 4, 4, 3]) + dev = qml.device("default.mixed", wires=1, shots=[1, 3, 3, 4, 4, 4, 3]) shot_vector = dev.shot_vector assert len(shot_vector) == 4 assert shot_vector[0].shots == 1 @@ -1069,6 +1086,15 @@ def test_shot_vector_property(self): assert dev.shots == 22 + def test_decomp_depth_is_deprecated(self): + """Test that a warning is raised when using the deprecated decomp_depth argument""" + + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The decomp_depth argument is deprecated", + ): + qml.device("default.qubit", decomp_depth=1) + class TestBatchExecution: """Tests for the batch_execute method.""" diff --git a/tests/devices/test_preprocess.py b/tests/devices/test_preprocess.py index cdfeec6c201..7cdbddb56d2 100644 --- a/tests/devices/test_preprocess.py +++ b/tests/devices/test_preprocess.py @@ -199,7 +199,7 @@ class TestDecomposeValidation: """Unit tests for helper functions in qml.devices.qubit.preprocess""" def test_error_if_invalid_op(self): - """Test that expand_fn throws an error when an operation is does not define a matrix or decomposition.""" + """Test that expand_fn throws an error when an operation does not define a matrix or decomposition.""" tape = QuantumScript(ops=[NoMatNoDecompOp(0)], measurements=[qml.expval(qml.Hadamard(0))]) with pytest.raises(DeviceError, match="not supported with abc"): diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py index 83e153b4401..05d5efeb437 100644 --- a/tests/drawer/test_draw.py +++ b/tests/drawer/test_draw.py @@ -878,13 +878,11 @@ def circ(): class TestLevelExpansionStrategy: - @pytest.fixture( - params=[qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit()], - ) - def transforms_circuit(self, request): + @pytest.fixture + def transforms_circuit(self): @qml.transforms.merge_rotations @qml.transforms.cancel_inverses - @qml.qnode(request.param, diff_method="parameter-shift") + @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift") def circ(weights, order): qml.RandomLayers(weights, wires=(0, 1)) qml.Permute(order, wires=(0, 1, 2)) @@ -953,8 +951,8 @@ def test_draw_at_level_1(self, transforms_circuit): ) assert out == expected - def test_draw_with_qfunc_warns_with_expansion_strategy(self): - """Test that draw warns the user about expansion_strategy being ignored.""" + def test_draw_with_qfunc_warns_with_expansion_strategy_or_level(self): + """Test that draw warns the user about expansion_strategy and level being ignored.""" def qfunc(): qml.PauliZ(0) @@ -962,17 +960,26 @@ def qfunc(): with pytest.warns( UserWarning, match="the expansion_strategy and level arguments are ignored" ): - _ = qml.draw(qfunc, expansion_strategy="gradient") + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + qml.draw(qfunc, expansion_strategy="gradient") with pytest.warns( UserWarning, match="the expansion_strategy and level arguments are ignored" ): - _ = qml.draw(qfunc, level="gradient") + qml.draw(qfunc, level="gradient") def test_providing_both_level_and_expansion_raises_error(self, transforms_circuit): with pytest.raises(ValueError, match="Either 'level' or 'expansion_strategy'"): qml.draw(transforms_circuit, level=0, expansion_strategy="device") + def test_deprecation_warning_when_expansion_strategy_provided(self, transforms_circuit): + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + qml.draw(transforms_circuit, expansion_strategy="device") + def test_draw_batch_transform(): """Test that drawing a batch transform works correctly.""" @@ -1015,11 +1022,7 @@ def circ(): assert draw(circ)() == expected -@pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)], -) -def test_applied_transforms(device): +def test_applied_transforms(): """Test that any transforms applied to the qnode are included in the output.""" @qml.transform @@ -1028,7 +1031,7 @@ def just_pauli_x(_): return (new_tape,), lambda res: res[0] @just_pauli_x - @qml.qnode(device) + @qml.qnode(qml.device("default.qubit", wires=2)) def my_circuit(x): qml.RX(x, wires=0) qml.SWAP(wires=(0, 1)) diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py index 2a8d2c83330..c5ed45534d0 100644 --- a/tests/drawer/test_draw_mpl.py +++ b/tests/drawer/test_draw_mpl.py @@ -81,13 +81,11 @@ def test_fig_argument(): class TestLevelExpansionStrategy: - @pytest.fixture( - params=[qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit()], - ) - def transforms_circuit(self, request): + @pytest.fixture + def transforms_circuit(self): @qml.transforms.merge_rotations @qml.transforms.cancel_inverses - @qml.qnode(request.param, diff_method="parameter-shift") + @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift") def circ(weights, order): qml.RandomLayers(weights, wires=(0, 1)) qml.Permute(order, wires=(0, 1, 2)) @@ -109,7 +107,7 @@ def circ(weights, order): ], ) def test_equivalent_levels(self, transforms_circuit, levels, expected_metadata): - """Test that the expansion strategy keyword controls what operations are drawn.""" + """Test that the level keyword controls what operations are drawn.""" var1, var2 = levels expected_lines, expected_patches, expected_texts = expected_metadata @@ -125,23 +123,27 @@ def test_equivalent_levels(self, transforms_circuit, levels, expected_metadata): plt.close("all") - @pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit(wires=3)], - ) @pytest.mark.parametrize( "strategy, initial_strategy, n_lines", [("gradient", "device", 3), ("device", "gradient", 13)], ) - def test_expansion_strategy(self, device, strategy, initial_strategy, n_lines): + def test_expansion_strategy(self, strategy, initial_strategy, n_lines): """Test that the expansion strategy keyword controls what operations are drawn.""" - @qml.qnode(device, expansion_strategy=initial_strategy) - def circuit(): - qml.Permute([2, 0, 1], wires=(0, 1, 2)) - return qml.expval(qml.PauliZ(0)) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="'expansion_strategy' attribute is deprecated", + ): - _, ax = qml.draw_mpl(circuit, expansion_strategy=strategy)() + @qml.qnode(qml.device("default.qubit"), expansion_strategy=initial_strategy) + def circuit(): + qml.Permute([2, 0, 1], wires=(0, 1, 2)) + return qml.expval(qml.PauliZ(0)) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + _, ax = qml.draw_mpl(circuit, expansion_strategy=strategy)() assert len(ax.lines) == n_lines assert circuit.expansion_strategy == initial_strategy @@ -172,7 +174,10 @@ def qfunc(): with pytest.warns( UserWarning, match="the expansion_strategy and level arguments are ignored" ): - qml.draw_mpl(qfunc, expansion_strategy="gradient") + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + qml.draw_mpl(qfunc, expansion_strategy="gradient") with pytest.warns( UserWarning, match="the expansion_strategy and level arguments are ignored" @@ -446,7 +451,13 @@ def qfunc(): qml.PauliZ(0) with pytest.warns(UserWarning, match="the expansion_strategy and level arguments are ignored"): - _ = qml.draw_mpl(qfunc, expansion_strategy="gradient") + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + qml.draw_mpl(qfunc, expansion_strategy="gradient") + + with pytest.warns(UserWarning, match="the expansion_strategy and level arguments are ignored"): + qml.draw_mpl(qfunc, level="gradient") def test_qnode_mid_circuit_measurement_not_deferred_device_api(mocker): @@ -499,11 +510,7 @@ def circuit(): assert ax.texts[-1].get_text() == "X†" -@pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)], -) -def test_applied_transforms(device): +def test_applied_transforms(): """Test that any transforms applied to the qnode are included in the output.""" @qml.transform @@ -512,7 +519,7 @@ def just_pauli_x(_): return (new_tape,), lambda res: res[0] @just_pauli_x - @qml.qnode(device) + @qml.qnode(qml.device("default.qubit", wires=2)) def my_circuit(): qml.SWAP(wires=(0, 1)) qml.CNOT(wires=(0, 1)) diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py index c8568c6293f..60d1b997004 100644 --- a/tests/drawer/test_tape_mpl.py +++ b/tests/drawer/test_tape_mpl.py @@ -844,7 +844,7 @@ def test_single_measure_multiple_conds(self): qml.cond(m0, qml.PauliY)(0) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") assert len(ax.patches) == 5 # three for measure, two for boxes @@ -857,11 +857,11 @@ def test_single_measure_multiple_conds(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } plt.close() @@ -875,7 +875,7 @@ def test_combo_measurement(self): qml.cond(m0 & m1, qml.PauliY)(0) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") assert len(ax.patches) == 7 # three for 2 measurements, one for box [_, _, cwire1, cwire2, eraser] = ax.lines @@ -891,18 +891,18 @@ def test_combo_measurement(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } assert eraser.get_xdata() == (1.8, 2) assert eraser.get_ydata() == (2, 2) - assert eraser.get_color() == plt.rcParams["figure.facecolor"] - assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"] + assert eraser.get_color() == "white" # hardcoded value to black_white color + assert eraser.get_linewidth() == 3 * 1.5 # hardcoded value to black_white linewidth plt.close() @@ -918,7 +918,7 @@ def test_combo_measurement_non_terminal(self): qml.cond(m1, qml.T)(1) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") [_, _, cwire1, cwire2, eraser] = ax.lines @@ -933,18 +933,18 @@ def test_combo_measurement_non_terminal(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } assert eraser.get_xdata() == (1.8, 2.2) assert eraser.get_ydata() == (2, 2) - assert eraser.get_color() == plt.rcParams["figure.facecolor"] - assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"] + assert eraser.get_color() == "white" # hardcoded value to black_white color + assert eraser.get_linewidth() == 3 * 1.5 # hardcoded value to black_white linewidth plt.close() @@ -954,7 +954,7 @@ def test_single_mcm_measure(self): with qml.queuing.AnnotatedQueue() as q: m0 = qml.measure(0) qml.expval(m0) - _, ax = tape_mpl(qml.tape.QuantumScript.from_queue(q)) + _, ax = tape_mpl(qml.tape.QuantumScript.from_queue(q), style="black_white") assert len(ax.patches) == 6 # two measurement boxes assert ax.patches[3].get_x() == 1 - 0.75 / 2 + 0.2 # 1 - box_length/2 + pad @@ -974,11 +974,11 @@ def test_single_mcm_measure(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } diff --git a/tests/fermi/test_fermionic.py b/tests/fermi/test_fermionic.py index b5a9b237fcc..3e0c324e4a4 100644 --- a/tests/fermi/test_fermionic.py +++ b/tests/fermi/test_fermionic.py @@ -125,6 +125,22 @@ def test_init_error(self, operator): with pytest.raises(ValueError, match="The operator indices must belong to the set"): FermiWord(operator) + def test_to_mat(self): + """Test that the matrix representation of FermiWord is correct.""" + + expected_mat = np.zeros((4, 4), dtype=complex) + expected_mat[2, 1] = 1.0 + + mat = fw1.to_mat() + assert np.allclose(mat, expected_mat) + + def test_to_mat_error(self): + """Test that an error is raised if the requested matrix dimension is smaller than the + dimension inferred from the largest orbital index. + """ + with pytest.raises(ValueError, match="n_orbitals cannot be smaller than 2"): + fw1.to_mat(n_orbitals=1) + class TestFermiWordArithmetic: WORDS_MUL = ( @@ -442,6 +458,14 @@ def test_array_must_not_exceed_length_1(self, method_name): fs3 = FermiSentence({fw3: -0.5, fw4: 1}) fs4 = FermiSentence({fw4: 1}) fs5 = FermiSentence({}) +fs6 = FermiSentence({fw1: 1.2, fw2: 3.1}) +fs7 = FermiSentence( + { + FermiWord({(0, 0): "+", (1, 1): "-"}): 1.23, # a+(0) a(1) + FermiWord({(0, 0): "+", (1, 0): "-"}): 4.0j, # a+(0) a(0) = n(0) (number operator) + FermiWord({(0, 0): "+", (1, 2): "-", (2, 1): "+"}): -0.5, # a+(0) a(2) a+(1) + } +) fs1_x_fs2 = FermiSentence( # fs1 * fs1, computed by hand { @@ -600,6 +624,25 @@ def test_pickling(self): new_fs = pickle.loads(serialization) assert fs == new_fs + def test_to_mat(self): + """Test that the matrix representation of FermiSentence is correct.""" + expected_mat = np.zeros((8, 8), dtype=complex) + expected_mat[4, 2] = 1.23 + 0j + expected_mat[5, 3] = 1.23 + 0j + for i in [4, 5, 6, 7]: + expected_mat[i, i] = 4.0j + expected_mat[6, 1] = 0.5 + 0j + + mat = fs7.to_mat() + assert np.allclose(mat, expected_mat) + + def test_to_mat_error(self): + """Test that an error is raised if the requested matrix dimension is smaller than the + dimension inferred from the largest orbital index. + """ + with pytest.raises(ValueError, match="n_orbitals cannot be smaller than 3"): + fs7.to_mat(n_orbitals=2) + class TestFermiSentenceArithmetic: tup_fs_mult = ( # computed by hand diff --git a/tests/gpu/test_gpu_torch.py b/tests/gpu/test_gpu_torch.py index cd461140e17..8c76734e683 100644 --- a/tests/gpu/test_gpu_torch.py +++ b/tests/gpu/test_gpu_torch.py @@ -33,7 +33,7 @@ class TestTorchDevice: def test_device_to_cuda(self): """Checks device executes with cuda is input data is cuda""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) @@ -53,7 +53,7 @@ def test_device_to_cuda(self): def test_mixed_devices(self): """Asserts works with both cuda and cpu input data""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) y = torch.tensor(0.2, requires_grad=True, device=torch.device("cpu")) @@ -77,7 +77,7 @@ def test_mixed_devices(self): def test_matrix_input(self): """Test goes to GPU for matrix valued inputs.""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) U = torch.eye(2, requires_grad=False, device=torch.device("cuda")) @@ -93,7 +93,7 @@ def test_matrix_input(self): def test_resets(self): """Asserts reverts to cpu after execution on gpu""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) y = torch.tensor(0.2, requires_grad=True, device=torch.device("cpu")) @@ -143,7 +143,7 @@ def test_different_devices_creation_and_parameters_warn(self, init_device, par_d PennyLane device creation differs from the Torch device of gate parameters. """ - dev = qml.device("default.qubit.torch", wires=1, torch_device=init_device) + dev = qml.device("default.qubit", wires=1, torch_device=init_device) p = torch.tensor(0.543, dtype=torch.float64, device=par_device) @@ -180,7 +180,7 @@ def circuit_cuda(inputs): @pytest.mark.skipif(not torch_cuda.is_available(), reason="no cuda support") -class TestqnnTorchLayer: +class TestQnnTorchLayer: def test_torch_device_cuda_if_tensors_on_cuda(self): """Test that if any tensor passed to operators is on the GPU then CUDA is set internally as a device option for 'default.qubit.torch'.""" diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py index 44b2ade9e6a..0a3ccb92ae9 100644 --- a/tests/gradients/core/test_gradient_transform.py +++ b/tests/gradients/core/test_gradient_transform.py @@ -619,10 +619,14 @@ def test_template_integration(self, strategy, tol): correctly when the QNode contains a template""" dev = qml.device("default.qubit", wires=3) - @qml.qnode(dev, expansion_strategy=strategy) - def circuit(weights): - qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2]) - return qml.probs(wires=[0, 1]) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' attribute is deprecated" + ): + + @qml.qnode(dev, expansion_strategy=strategy) + def circuit(weights): + qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2]) + return qml.probs(wires=[0, 1]) weights = np.ones([2, 3, 3], dtype=np.float64, requires_grad=True) res = qml.gradients.param_shift(circuit)(weights) diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py index a455824e1df..682f685b99f 100644 --- a/tests/gradients/core/test_hadamard_gradient.py +++ b/tests/gradients/core/test_hadamard_gradient.py @@ -875,61 +875,21 @@ def circuit(weights): with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): qml.gradients.hadamard_grad(circuit)(weights) - @pytest.mark.autograd - def test_no_trainable_params_qnode_autograd_legacy_opmath(self): + @pytest.mark.parametrize( + "interface", + [ + pytest.param("jax", marks=pytest.mark.jax), + pytest.param("autograd", marks=pytest.mark.autograd), + pytest.param("torch", marks=pytest.mark.torch), + pytest.param("tf", marks=pytest.mark.tf), + ], + ) + def test_no_trainable_params_qnode_legacy_opmath(self, interface): """Test that the correct ouput and warning is generated in the absence of any trainable parameters""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, interface="autograd") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) + dev = qml.device(f"default.qubit.{interface}", wires=2) - @pytest.mark.torch - def test_no_trainable_params_qnode_torch_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.torch", wires=2) - - @qml.qnode(dev, interface="torch") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) - - @pytest.mark.tf - def test_no_trainable_params_qnode_tf_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.tf", wires=2) - - @qml.qnode(dev, interface="tf") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) - - @pytest.mark.jax - def test_no_trainable_params_qnode_jax_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, interface="jax") + @qml.qnode(dev, interface=interface) def circuit(weights): qml.RX(weights[0], wires=0) qml.RY(weights[1], wires=0) diff --git a/tests/gradients/core/test_metric_tensor.py b/tests/gradients/core/test_metric_tensor.py index 01b43bc8177..364516d6079 100644 --- a/tests/gradients/core/test_metric_tensor.py +++ b/tests/gradients/core/test_metric_tensor.py @@ -14,7 +14,9 @@ """ Unit tests for the metric tensor transform. """ -# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods,not-callable +import importlib + +# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods,not-callable,too-many-statements import pytest from scipy.linalg import block_diag @@ -538,189 +540,31 @@ def layer3_diag(params): G_expected = block_diag(G1, G3, G2) assert qml.math.allclose(G, G_expected, atol=tol, rtol=0) - @pytest.mark.autograd - @pytest.mark.parametrize("interface", ["autograd"]) - def test_argnum_metric_tensor_autograd(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - dev = qml.device("default.qubit.autograd", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, np.ndarray) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, np.ndarray) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, np.ndarray) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) - - @pytest.mark.tf - @pytest.mark.parametrize("interface", ["tf"]) - def test_argnum_metric_tensor_tf(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - import tensorflow as tf - - dev = qml.device("default.qubit.tf", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = tf.Variable([0.1, 0.2, 0.3, 0.5]) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, tf.Tensor) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, tf.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, tf.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) - - @pytest.mark.torch - @pytest.mark.parametrize("interface", ["torch"]) - def test_argnum_metric_tensor_torch(self, tol, interface): + @pytest.mark.parametrize( + "interface,array_cls", + [ + pytest.param("jax", "array", marks=pytest.mark.jax), + pytest.param("autograd", "array", marks=pytest.mark.autograd), + pytest.param("tf", "Variable", marks=pytest.mark.tf), + pytest.param("torch", "Tensor", marks=pytest.mark.torch), + ], + ) + def test_argnum_metric_tensor_interfaces(self, tol, interface, array_cls): """Test that argnum successfully reduces the number of tapes and gives the desired outcome.""" - import torch - - dev = qml.device("default.qubit.torch", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = torch.tensor([0.1, 0.2, 0.3, 0.5], requires_grad=True) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, torch.Tensor) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, torch.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, torch.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) + if interface == "tf": + interface_name = "tensorflow" + elif interface == "jax": + interface_name = "jax.numpy" + elif interface == "autograd": + interface_name = "numpy" + else: + interface_name = interface - @pytest.mark.jax - @pytest.mark.parametrize("interface", ["jax"]) - def test_argnum_metric_tensor_jax(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - import jax + mod = importlib.import_module(interface_name) + type_ = type(getattr(mod, array_cls)([])) if interface != "tf" else getattr(mod, "Tensor") - dev = qml.device("default.qubit.jax", wires=3) + dev = qml.device("default.qubit", wires=3) def circuit(weights): qml.RX(weights[0], wires=0) @@ -729,7 +573,7 @@ def circuit(weights): qml.RZ(weights[2], wires=1) qml.RZ(weights[3], wires=0) - weights = jax.numpy.array([0.1, 0.2, 0.3, 0.5]) + weights = getattr(mod, array_cls)([0.1, 0.2, 0.3, 0.5]) with qml.tape.QuantumTape() as tape: circuit(weights) @@ -741,7 +585,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) res = qml.execute(tapes, dev, None, interface=interface) mt013 = proc_fn(res) - assert isinstance(mt013, jax.numpy.ndarray) + assert isinstance(mt013, type_) assert len(tapes) == 6 assert mt.shape == mt013.shape @@ -753,7 +597,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) res = qml.execute(tapes, dev, None, interface=interface) mt23 = proc_fn(res) - assert isinstance(mt23, jax.numpy.ndarray) + assert isinstance(mt23, type_) assert len(tapes) == 1 assert mt.shape == mt23.shape @@ -764,7 +608,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=0) res = qml.execute(tapes, dev, None, interface=interface) mt0 = proc_fn(res) - assert isinstance(mt0, jax.numpy.ndarray) + assert isinstance(mt0, type_) assert len(tapes) == 1 assert mt.shape == mt0.shape diff --git a/tests/helpers/mcm_utils.py b/tests/helpers/mcm_utils.py index fab28c2bfca..816991981cd 100644 --- a/tests/helpers/mcm_utils.py +++ b/tests/helpers/mcm_utils.py @@ -14,8 +14,8 @@ """ Pytest helper functions are defined in this module. """ +from collections.abc import Iterable, Sequence from functools import reduce -from typing import Iterable, Sequence import numpy as np diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py deleted file mode 100644 index 2a6ee306508..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py +++ /dev/null @@ -1,839 +0,0 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Autograd specific tests for execute and default qubit 2.""" -import autograd -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane import numpy as np -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -pytestmark = pytest.mark.autograd - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = np.arange(1, num_params + 1) / 10 - - N = len(params) - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = np.array( - [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - def test_single_backward_pass_batch(self): - """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift - is requested for multiple tapes.""" - - dev = qml.device("default.qubit") - - def f(x): - tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) - tape2 = qml.tape.QuantumScript([qml.RY(x, 0)], [qml.probs(wires=0)]) - - results = qml.execute([tape1, tape2], dev, gradient_fn=qml.gradients.param_shift) - return results[0] + results[1] - - x = qml.numpy.array(0.1) - with dev.tracker: - out = qml.jacobian(f)(x) - - assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["simulations"] == [1, 1, 1, 1, 1, 1] - expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] - assert qml.math.allclose(out, expected) - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), - ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, - Shots(None), - DefaultQubit(), - ), - ( - { - "gradient_fn": "adjoint", - "grad_on_execution": False, - "device_vjp": False, - }, - Shots(None), - DefaultQubit(), - ), - ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "device", "device_vjp": False}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(seed=904747894), - ), - ( - {"gradient_fn": "device", "device_vjp": True}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(seed=10490244), - ), -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestAutogradExecuteIntegration: - """Test the autograd interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=False) - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("grad_on_execution", False): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = np.array(0.1, requires_grad=True) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - if shots.has_partitioned_shots: - res = qml.jacobian(lambda x: qml.math.hstack(cost(x)))(a) - else: - res = qml.jacobian(cost)(a) - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -np.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = qml.jacobian(cost)(a, b) - assert isinstance(res, tuple) and len(res) == 2 - if shots.has_partitioned_shots: - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - for _r, _e in zip(res, expected): - assert np.allclose(_r[:2], _e, atol=atol_for_shots(shots)) - assert np.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - else: - assert res[0].shape == (2,) - assert res[1].shape == (2,) - - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - for _r, _e in zip(res, expected): - assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - - @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - res = qml.execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = tuple(i for r in res for i in r) - return sum(autograd.numpy.hstack(res)) - - params = np.array([0.1, 0.2], requires_grad=True) - x, y = params - - res = cost(params) - expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - if shots.has_partitioned_shots: - expected = shots.num_copies * expected - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - grad = qml.grad(cost)(params) - expected = np.array([-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]) - if shots.has_partitioned_shots: - expected = shots.num_copies * expected - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if execute_kwargs["gradient_fn"] == "backprop": - pytest.xfail("backprop is not compatible with something about this situation.") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = execute([tape1, tape2, tape3], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = tuple(i for r in res for i in r) - return autograd.numpy.hstack(res) - - params = np.array([0.1, 0.2], requires_grad=True) - x, y = params - - res = cost(params) - assert isinstance(res, np.ndarray) - if not shots: - assert res.shape == (4,) - - if shots.has_partitioned_shots: - for i in (0, 1): - assert np.allclose(res[2 * i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[2 * i + 1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[4 + i], np.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[6 + i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], np.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - - if shots.has_partitioned_shots: - pytest.xfail("autograd jacobians do not work with ragged results and shot vectors.") - # TODO: autograd jacobians with ragged results and shot vectors - jac = qml.jacobian(cost)(params) - assert isinstance(jac, np.ndarray) - if not shots.has_partitioned_shots: - assert jac.shape == (4, 2) # pylint: disable=no-member - - d1 = -np.sin(x) * np.cos(y) - d2 = -np.cos(x) * np.sin(y) - - if shots.has_partitioned_shots: - assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) - assert np.allclose(jac[3:4], 0, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) - - else: - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return autograd.numpy.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac_fn = qml.jacobian(cost) - jac = jac_fn(a, b) - - a = np.array(0.54, requires_grad=True) - b = np.array(0.8, requires_grad=True) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) - jac = jac_fn(a, b) - expected = ( - [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)], - [0, -np.cos(2 * a) * np.cos(b)], - ) - assert isinstance(jac, tuple) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=False) - c = np.array(0.3, requires_grad=True) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + np.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - if shots.has_partitioned_shots: - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0]) - return execute([tape], device, **execute_kwargs)[0] - - res = qml.jacobian(cost)(a, b, c) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = np.array(0.1, requires_grad=False) - b = np.array(0.2, requires_grad=False) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - assert res.shape == (2 * shots.num_copies,) if shots else (2,) - - with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.jacobian(cost)(a, b) - assert len(res) == 0 - - def loss(a, b): - return np.sum(cost(a, b)) - - with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.grad(loss)(a, b) - - assert np.allclose(res, 0) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the autograd interface works correctly - with a matrix parameter""" - U = np.array([[0, 1], [1, 0]], requires_grad=False) - a = np.array(0.1, requires_grad=True) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert np.allclose(res, -np.cos(a), atol=atol_for_shots(shots)) - - jac_fn = qml.jacobian(cost) - jac = jac_fn(a, U) - assert isinstance(jac, np.ndarray) - assert np.allclose(jac, np.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = np.array(0.1, requires_grad=False) - p = np.array([0.1, 0.2, 0.3], requires_grad=True) - - res = cost_fn(a, p) - expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( - np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost_fn) - res = jac_fn(a, p) - expected = np.array( - [ - np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), - np.cos(p[1]) * np.cos(p[2]) * np.sin(a) - - np.sin(p[1]) - * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), - np.sin(a) - * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = np.array(0.543, requires_grad=True) - y = np.array(-0.654, requires_grad=True) - - res = cost(x, y) - expected = np.array( - [ - [ - np.cos(x / 2) ** 2, - np.sin(x / 2) ** 2, - (1 + np.cos(x) * np.cos(y)) / 2, - (1 - np.cos(x) * np.cos(y)) / 2, - ], - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - np.array( - [ - [ - -np.sin(x) / 2, - np.sin(x) / 2, - -np.sin(x) * np.cos(y) / 2, - np.sin(x) * np.cos(y) / 2, - ], - ] - ), - np.array( - [ - [0, 0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ] - ), - ) - - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = np.array(0.543, requires_grad=True) - y = np.array(-0.654, requires_grad=True) - - res = cost(x, y) - expected = np.array( - [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), - np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), - ) - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the autograd execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - np.array([0.543, -0.654], requires_grad=True), - np.array([0, -0.654], requires_grad=True), - np.array([-2.0, 0], requires_grad=True), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using autograd, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.grad(cost_fn)(params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.jacobian(qml.grad(cost_fn))(params) - expected = np.array( - [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.grad(cost_fn)(params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.grad(cost_fn))(params) - - expected = np.zeros([2, 2]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return np.array( - [ - [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), - 0, - ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) - coeffs2 = np.array([0.7], requires_grad=False) - weights = np.array([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) - coeffs2 = np.array([0.7], requires_grad=True) - weights = np.array([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2)) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py deleted file mode 100644 index 7891f43ea77..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tests for exeuction with default qubit 2 independent of any interface.""" -import pytest - -import pennylane as qml -from pennylane import numpy as np -from pennylane.devices import DefaultQubit -from pennylane.workflow.execution import _preprocess_expand_fn - - -class TestPreprocessExpandFn: - """Tests the _preprocess_expand_fn helper function.""" - - def test_provided_is_callable(self): - """Test that if the expand_fn is not "device", it is simply returned.""" - - dev = DefaultQubit() - - def f(tape): - return tape - - out = _preprocess_expand_fn(f, dev, 10) - assert out is f - - def test_new_device_blank_expand_fn(self): - """Test that the expand_fn is blank if is new device.""" - - dev = DefaultQubit() - - out = _preprocess_expand_fn("device", dev, 10) - - x = [1] - assert out(x) is x - - -class TestBatchTransformHelper: - """Unit tests for the _batch_transform helper function.""" - - def test_warns_if_requested_off(self): - """Test that a warning is raised if the the batch transform is requested to not be used.""" - - # pylint: disable=too-few-public-methods - class CustomOp(qml.operation.Operator): - """Dummy operator.""" - - def decomposition(self): - return [qml.PauliX(self.wires[0])] - - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): - program, _ = dev.preprocess() - qml.execute( - (qs, qs), device=dev, device_batch_transform=False, transform_program=program - ) - - def test_split_and_expand_performed(self): - """Test that preprocess returns the correct tapes when splitting and expanding - is needed.""" - - class NoMatOp(qml.operation.Operation): - """Dummy operation for expanding circuit.""" - - # pylint: disable=missing-function-docstring - num_wires = 1 - - # pylint: disable=arguments-renamed, invalid-overridden-method - @property - def has_matrix(self): - return False - - def decomposition(self): - return [qml.PauliX(self.wires), qml.PauliY(self.wires)] - - ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)] - # Need to specify grouping type to transform tape - measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))] - tapes = [ - qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]), - qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]), - ] - - dev = DefaultQubit() - config = qml.devices.ExecutionConfig(gradient_method="adjoint") - - program, new_config = dev.preprocess(config) - res_tapes, batch_fn = program(tapes) - expected_ops = [ - [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)], - [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)], - ] - - assert len(res_tapes) == 4 - for i, t in enumerate(res_tapes): - for op, expected_op in zip(t.operations, expected_ops[i % 2]): - qml.assert_equal(op, expected_op) - assert len(t.measurements) == 1 - if i < 2: - qml.assert_equal(t.measurements[0], measurements[0]) - else: - qml.assert_equal(t.measurements[0], measurements[1]) - - input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]) - assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])) - - assert new_config.grad_on_execution - assert new_config.use_device_gradient - - -def test_warning_if_not_device_batch_transform(): - """Test that a warning is raised if the users requests to not run device batch transform.""" - - # pylint: disable=too-few-public-methods - class CustomOp(qml.operation.Operator): - """Dummy operator.""" - - def decomposition(self): - return [qml.PauliX(self.wires[0])] - - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): - program, _ = dev.preprocess() - results = qml.execute([qs], dev, device_batch_transform=False, transform_program=program) - - assert len(results) == 1 - assert qml.math.allclose(results[0], -1) - - -@pytest.mark.parametrize("gradient_fn", (None, "backprop", qml.gradients.param_shift)) -def test_caching(gradient_fn): - """Test that cache execute returns the cached result if the same script is executed - multiple times, both in multiple times in a batch and in separate batches.""" - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - cache = {} - - with qml.Tracker(dev) as tracker: - results = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) - results2 = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) - - assert len(cache) == 1 - assert cache[qs.hash] == -1.0 - - assert list(results) == [-1.0, -1.0] - assert list(results2) == [-1.0, -1.0] - - assert tracker.totals["batches"] == 1 - assert tracker.totals["executions"] == 1 - assert cache[qs.hash] == -1.0 diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py deleted file mode 100644 index 519c0daa028..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py +++ /dev/null @@ -1,817 +0,0 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Jax specific tests for execute and default qubit 2.""" -import numpy as np -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -jax = pytest.importorskip("jax") -jnp = pytest.importorskip("jax.numpy") -jax.config.update("jax_enable_x64", True) - -pytestmark = pytest.mark.jax - - -def test_jit_execution(): - """Test that qml.execute can be directly jitted.""" - dev = qml.device("default.qubit") - - tape = qml.tape.QuantumScript( - [qml.RX(jax.numpy.array(0.1), 0)], [qml.expval(qml.s_prod(2.0, qml.PauliZ(0)))] - ) - - out = jax.jit(qml.execute, static_argnames=("device", "gradient_fn"))( - (tape,), device=dev, gradient_fn=qml.gradients.param_shift - ) - expected = 2.0 * jax.numpy.cos(jax.numpy.array(0.1)) - assert qml.math.allclose(out[0], expected) - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.skip("caching is not implemented for jax") - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - device = DefaultQubit() - params = jnp.arange(1, num_params + 1) / 10 - - N = len(params) - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], device, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(device) as tracker: - hess1 = jax.jacobian(jax.grad(cost))(params, cache=False) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = jnp.array( - [ - [2 * jnp.cos(2 * x) * jnp.sin(y) ** 2, jnp.sin(2 * x) * jnp.sin(2 * y)], - [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(device) as tracker2: - hess2 = jax.jacobian(jax.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -no_shots = Shots(None) -shots_2_10k = Shots((10000, 10000)) -dev_def = DefaultQubit() -dev_ps = ParamShiftDerivativesDevice(seed=54353453) -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), # 0 - ({"gradient_fn": param_shift}, no_shots, dev_def), # 1 - ({"gradient_fn": "backprop"}, no_shots, dev_def), # 2 - ({"gradient_fn": "adjoint"}, no_shots, dev_def), # 3 - ({"gradient_fn": "adjoint", "device_vjp": True}, no_shots, dev_def), # 4 - ({"gradient_fn": "device"}, shots_2_10k, dev_ps), # 5 -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 3e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestJaxExecuteIntegration: - """Test the jax interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = jnp.array(0.1) - b = np.array(0.2) - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("gradient_fn", None) == "adjoint": - assert device.tracker.totals.get("execute_and_derivative_batches", 0) == 0 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = jnp.array(0.1) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = jax.jacobian(cost)(a) - if not shots.has_partitioned_shots: - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -jnp.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - - a = jnp.array(0.1) - b = jnp.array(0.2) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, b) - expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res[0], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected, atol=2 * atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - g = jax.jacobian(cost, argnums=[0, 1])(a, b) - assert isinstance(g, tuple) and len(g) == 2 - - expected = ([-jnp.sin(a), jnp.sin(a) * jnp.sin(b)], [0, -jnp.cos(a) * jnp.cos(b)]) - - if shots.has_partitioned_shots: - for i in (0, 1): - assert np.allclose(g[i][0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(g[0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(jnp.array(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(jnp.array(0.5), 0)], [qml.probs(wires=[0, 1])], shots=shots - ) - res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - res = jax.tree_util.tree_leaves(res) - out = sum(jnp.hstack(res)) - if shots.has_partitioned_shots: - out = out / shots.num_copies - return out - - params = jnp.array([0.1, 0.2]) - - x, y = params - - res = cost(params) - expected = 2 + jnp.cos(0.5) + jnp.cos(x) * jnp.cos(y) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - grad = jax.grad(cost)(params) - expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)] - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - # pylint: disable=too-many-statements - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], [qml.expval(qml.PauliZ(0))], shots=shots - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = qml.execute([tape1, tape2, tape3], device, **execute_kwargs) - leaves = jax.tree_util.tree_leaves(res) - return jnp.hstack(leaves) - - params = jnp.array([0.1, 0.2]) - x, y = params - - res = cost(params) - assert isinstance(res, jax.Array) - assert res.shape == (4 * shots.num_copies,) if shots.has_partitioned_shots else (4,) - - if shots.has_partitioned_shots: - assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[2], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[3], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[4], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[5], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[6], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[7], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - - jac = jax.jacobian(cost)(params) - assert isinstance(jac, jnp.ndarray) - assert ( - jac.shape == (8, 2) if shots.has_partitioned_shots else (4, 2) - ) # pylint: disable=no-member - - if shots.has_partitioned_shots: - assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) - assert np.allclose(jac[3:5], 0, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - d1 = -jnp.sin(x) * jnp.cos(y) - if shots.has_partitioned_shots: - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -jnp.cos(x) * jnp.sin(y) - if shots.has_partitioned_shots: - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - if execute_kwargs["gradient_fn"] == param_shift: - pytest.skip("Basic QNode execution wipes out trainable params with param-shift") - - a = jnp.array(0.1) - b = jnp.array(0.2) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - shots=shots, - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return jnp.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - jac = jac_fn(a, b) - - a = jnp.array(0.54) - b = jnp.array(0.8) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [jnp.cos(2 * a), -jnp.cos(2 * a) * jnp.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res2[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res2[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(lambda a, b: cost(2 * a, b), argnums=[0, 1]) - jac = jac_fn(a, b) - expected = ( - [-2 * jnp.sin(2 * a), 2 * jnp.sin(2 * a) * jnp.sin(b)], - [0, -jnp.cos(2 * a) * jnp.cos(b)], - ) - assert isinstance(jac, tuple) and len(jac) == 2 - if shots.has_partitioned_shots: - for offset in (0, 2): - assert np.allclose(jac[0][0 + offset], expected[0][0], atol=atol_for_shots(shots)) - assert np.allclose(jac[0][1 + offset], expected[0][1], atol=atol_for_shots(shots)) - assert np.allclose(jac[1][0 + offset], expected[1][0], atol=atol_for_shots(shots)) - assert np.allclose(jac[1][1 + offset], expected[1][1], atol=atol_for_shots(shots)) - else: - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, shots, device): - """Test classical processing within the quantum tape""" - a = jnp.array(0.1) - b = jnp.array(0.2) - c = jnp.array(0.3) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + jnp.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = jax.jacobian(cost, argnums=[0, 2])(a, b, c) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the jax interface works correctly - with a matrix parameter""" - U = jnp.array([[0, 1], [1, 0]]) - a = jnp.array(0.1) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert np.allclose(res, -jnp.cos(a), atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost) - jac = jac_fn(a, U) - if not shots.has_partitioned_shots: - assert isinstance(jac, jnp.ndarray) - assert np.allclose(jac, jnp.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], - [qml.expval(qml.PauliX(0))], - shots=shots, - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - conf = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=conf) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = jnp.array(0.1) - p = jnp.array([0.1, 0.2, 0.3]) - - res = cost_fn(a, p) - expected = jnp.cos(a) * jnp.cos(p[1]) * jnp.sin(p[0]) + jnp.sin(a) * ( - jnp.cos(p[2]) * jnp.sin(p[1]) + jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.sin(p[2]) - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost_fn, argnums=[1]) - res = jac_fn(a, p) - expected = jnp.array( - [ - jnp.cos(p[1]) - * (jnp.cos(a) * jnp.cos(p[0]) - jnp.sin(a) * jnp.sin(p[0]) * jnp.sin(p[2])), - jnp.cos(p[1]) * jnp.cos(p[2]) * jnp.sin(a) - - jnp.sin(p[1]) - * (jnp.cos(a) * jnp.sin(p[0]) + jnp.cos(p[0]) * jnp.sin(a) * jnp.sin(p[2])), - jnp.sin(a) - * (jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.cos(p[2]) - jnp.sin(p[1]) * jnp.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return jnp.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = jnp.array(0.543) - y = jnp.array(-0.654) - - res = cost(x, y) - expected = jnp.array( - [ - [ - jnp.cos(x / 2) ** 2, - jnp.sin(x / 2) ** 2, - (1 + jnp.cos(x) * jnp.cos(y)) / 2, - (1 - jnp.cos(x) * jnp.cos(y)) / 2, - ], - ] - ) - if shots.has_partitioned_shots: - assert np.allclose(res[:, 0:2].flatten(), expected, atol=atol_for_shots(shots)) - assert np.allclose(res[:, 2:].flatten(), expected, atol=atol_for_shots(shots)) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,) - assert res[1].shape == (2, 4) if shots.has_partitioned_shots else (4,) - - expected = ( - jnp.array( - [ - [ - -jnp.sin(x) / 2, - jnp.sin(x) / 2, - -jnp.sin(x) * jnp.cos(y) / 2, - jnp.sin(x) * jnp.cos(y) / 2, - ], - ] - ), - jnp.array( - [ - [0, 0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2], - ] - ), - ) - - if shots.has_partitioned_shots: - assert np.allclose(res[0][:, 0:2].flatten(), expected[0], atol=atol_for_shots(shots)) - assert np.allclose(res[0][:, 2:].flatten(), expected[0], atol=atol_for_shots(shots)) - assert np.allclose(res[1][:, :2].flatten(), expected[1], atol=atol_for_shots(shots)) - assert np.allclose(res[1][:, 2:].flatten(), expected[1], atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - res = qml.execute([tape], device, **execute_kwargs)[0] - return jnp.hstack(jax.tree_util.tree_leaves(res)) - - x = jnp.array(0.543) - y = jnp.array(-0.654) - - res = cost(x, y) - expected = jnp.array( - [jnp.cos(x), (1 + jnp.cos(x) * jnp.cos(y)) / 2, (1 - jnp.cos(x) * jnp.cos(y)) / 2] - ) - if shots.has_partitioned_shots: - assert np.allclose(res[:3], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[3:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) - assert res[1].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) - - expected = ( - jnp.array([-jnp.sin(x), -jnp.sin(x) * jnp.cos(y) / 2, jnp.sin(x) * jnp.cos(y) / 2]), - jnp.array([0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2]), - ) - if shots.has_partitioned_shots: - assert np.allclose(res[0][:3], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[0][3:], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1][:3], expected[1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1][3:], expected[1], atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the jax execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - jnp.array([0.543, -0.654]), - jnp.array([0, -0.654]), - jnp.array([-2.0, 0]), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using jax, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(cost_fn)(params) - expected = jnp.array( - [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jacobian(jax.grad(cost_fn))(params) - expected = jnp.array( - [ - [-jnp.cos(2 * x) * jnp.cos(2 * y), jnp.sin(2 * x) * jnp.sin(2 * y)], - [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = jnp.array([0.543, -0.654]) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(cost_fn)(params) - expected = jnp.array( - [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jacobian(jax.grad(cost_fn))(params) - expected = jnp.zeros([2, 2]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - res = execute([tape], device, **execute_kwargs)[0] - if shots.has_partitioned_shots: - return jnp.hstack(res[0] + res[1]) - return jnp.hstack(res) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * jnp.sin(x) * jnp.sin(y) + jnp.cos(x) * (a + b * jnp.sin(y)), d * jnp.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return jnp.array( - [ - [ - -c * jnp.cos(x) * jnp.sin(y) - jnp.sin(x) * (a + b * jnp.sin(y)), - b * jnp.cos(x) * jnp.cos(y) - c * jnp.cos(y) * jnp.sin(x), - jnp.cos(x), - jnp.cos(x) * jnp.sin(y), - -(jnp.sin(x) * jnp.sin(y)), - 0, - ], - [-d * jnp.sin(x), 0, 0, 0, 0, jnp.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = jnp.array([0.1, 0.2, 0.3]) - coeffs2 = jnp.array([0.7]) - weights = jnp.array([0.4, 0.5]) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = jax.jacobian(cost_fn)(weights, coeffs1, coeffs2) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = jnp.array([0.1, 0.2, 0.3]) - coeffs2 = jnp.array([0.7]) - weights = jnp.array([0.4, 0.5]) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = jnp.hstack(jax.jacobian(cost_fn, argnums=[0, 1, 2])(weights, coeffs1, coeffs2)) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py deleted file mode 100644 index 11a54575cfa..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py +++ /dev/null @@ -1,844 +0,0 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tensorflow specific tests for execute and default qubit 2.""" -import numpy as np -import pytest - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift - -pytestmark = pytest.mark.tf -tf = pytest.importorskip("tensorflow") - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = tf.Variable(tf.range(1, num_params + 1) / 10) - - N = num_params - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost(params, cache=False) - grad = grad_tape.gradient(res, params) - hess1 = jac_tape.jacobian(grad, params) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = tf.convert_to_tensor( - [ - [2 * tf.cos(2 * x) * tf.sin(y) ** 2, tf.sin(2 * x) * tf.sin(2 * y)], - [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost(params, cache=True) - grad = grad_tape.gradient(res, params) - hess2 = jac_tape.jacobian(grad, params) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)), # 0 - ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()), # 1 - ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()), # 2 - ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()), # 3 - ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)), # 4 - ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()), # 5 - ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()), # 6 - ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()), # 7 - ({"gradient_fn": "adjoint", "interface": "tf", "device_vjp": True}, None, DefaultQubit()), # 8 -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestTensorflowExecuteIntegration: - """Test the tensorflow interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = tf.Variable(0.1, dtype="float64") - b = tf.constant(0.2, dtype="float64") - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("gradient_fn", None) == "adjoint" and not execute_kwargs.get( - "device_vjp", False - ): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = tf.Variable(0.1, dtype=tf.float64) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(a) - res = tape.jacobian(cost_res, a, experimental_use_pfor=not device_vjp) - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -tf.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = tf.Variable(0.1) - b = tf.Variable(0.2) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(a, b) - expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - assert isinstance(jac, list) and len(jac) == 2 - assert jac[0].shape == (2,) - assert jac[1].shape == (2,) - - expected = ([-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]) - for _r, _e in zip(jac, expected): - assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - return tf.reduce_sum( - qml.math.hstack( - execute([tape1, tape2, tape3, tape4], device, **execute_kwargs), - like="tensorflow", - ) - ) - - params = tf.Variable([0.1, 0.2]) - x, y = params - - with tf.GradientTape() as tape: - res = cost(params) - expected = 2 + tf.cos(0.5) + tf.cos(x) * tf.cos(y) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(NotImplementedError): - tape.gradient(res, params) - return - - grad = tape.gradient(res, params) - expected = [-tf.cos(y) * tf.sin(x), -tf.cos(x) * tf.sin(y)] - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if ( - execute_kwargs["gradient_fn"] == "adjoint" - and execute_kwargs["interface"] == "tf-autograph" - ): - pytest.skip("Cannot compute the jacobian with adjoint-differentation and tf-autograph") - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5, dtype=tf.float64), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - return qml.math.hstack( - execute([tape1, tape2, tape3], device, **execute_kwargs), like="tensorflow" - ) - - params = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = params - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(params) - - assert isinstance(res, tf.Tensor) - assert res.shape == (4,) - - assert np.allclose(res[0], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], tf.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) - - jac = tape.jacobian(res, params, experimental_use_pfor=not device_vjp) - assert isinstance(jac, tf.Tensor) - assert jac.shape == (4, 2) # pylint: disable=no-member - - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - d1 = -tf.sin(x) * tf.cos(y) - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -tf.cos(x) * tf.sin(y) - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = tf.Variable(0.1) - b = tf.Variable(0.2) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return qml.math.hstack( - execute([new_tape], device, **execute_kwargs)[0], like="tensorflow" - ) - - with tf.GradientTape(persistent=device_vjp) as t: - res = cost(a, b) - - jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - a = tf.Variable(0.54, dtype=tf.float64) - b = tf.Variable(0.8, dtype=tf.float64) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - - with tf.GradientTape(persistent=device_vjp): - res2 = cost(2 * a, b) - - expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - with tf.GradientTape(persistent=device_vjp) as t: - res = cost(2 * a, b) - - jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - expected = ( - [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], - [0, -tf.cos(2 * a) * tf.cos(b)], - ) - assert isinstance(jac, list) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.constant(0.2, dtype=tf.float64) - c = tf.Variable(0.3, dtype=tf.float64) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + tf.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(a, b, c) - - res = tape.jacobian(cost_res, [a, c], experimental_use_pfor=not device_vjp) - - # Only two arguments are trainable - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = tf.constant(0.1) - b = tf.constant(0.2) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - with tf.GradientTape() as tape: - cost_res = cost(a, b) - - assert cost_res.shape == (2,) - - res = tape.jacobian(cost_res, [a, b]) - assert len(res) == 2 - assert all(r is None for r in res) - - def loss(a, b): - return tf.reduce_sum(cost(a, b)) - - with tf.GradientTape() as tape: - loss_res = loss(a, b) - - res = tape.gradient(loss_res, [a, b]) - assert all(r is None for r in res) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the tensorflow interface works correctly - with a matrix parameter""" - U = tf.constant([[0, 1], [1, 0]], dtype=tf.complex128) - a = tf.Variable(0.1) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(a, U) - - assert np.allclose(res, -tf.cos(a), atol=atol_for_shots(shots)) - - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) - assert isinstance(jac, tf.Tensor) - assert np.allclose(jac, tf.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - device_vjp = execute_kwargs.get("device_vjp", False) - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = tf.constant(0.1) - p = tf.Variable([0.1, 0.2, 0.3]) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost_fn(a, p) - - expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( - tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - res = tape.jacobian(cost_res, p, experimental_use_pfor=not device_vjp) - expected = tf.convert_to_tensor( - [ - tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), - tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) - - tf.sin(p[1]) - * (tf.cos(a) * tf.sin(p[0]) + tf.cos(p[0]) * tf.sin(a) * tf.sin(p[2])), - tf.sin(a) - * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - device_vjp = execute_kwargs.get("device_vjp", False) - - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(x, y) - - expected = tf.convert_to_tensor( - [ - [ - tf.cos(x / 2) ** 2, - tf.sin(x / 2) ** 2, - (1 + tf.cos(x) * tf.cos(y)) / 2, - (1 - tf.cos(x) * tf.cos(y)) / 2, - ], - ] - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(tf.errors.UnimplementedError): - tape.jacobian(cost_res, [x, y]) - return - res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - tf.convert_to_tensor( - [ - [ - -tf.sin(x) / 2, - tf.sin(x) / 2, - -tf.sin(x) * tf.cos(y) / 2, - tf.sin(x) * tf.cos(y) / 2, - ], - ] - ), - tf.convert_to_tensor( - [ - [0, 0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ] - ), - ) - - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(x, y) - - expected = tf.convert_to_tensor( - [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(tf.errors.UnimplementedError): - tape.jacobian(cost_res, [x, y]) - return - res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - tf.convert_to_tensor( - [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.sin(x) * tf.cos(y) / 2] - ), - tf.convert_to_tensor([0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2]), - ) - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the tensorflow execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - tf.Variable([0.543, -0.654], dtype=tf.float64), - tf.Variable([0, -0.654], dtype=tf.float64), - tf.Variable([-2.0, 0], dtype=tf.float64), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using tensorflow, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost_fn(params) - grad = grad_tape.gradient(res, params) - hess = jac_tape.jacobian(grad, params) - - x, y = params - expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [ - [-tf.cos(2 * x) * tf.cos(2 * y), tf.sin(2 * x) * tf.sin(2 * y)], - [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], - ] - ) - assert np.allclose(hess, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = tf.Variable([0.543, -0.654]) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost_fn(params) - grad = grad_tape.gradient(res, params) - hess = jac_tape.gradient(grad, params) - - x, y = params - expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) - assert hess is None - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * tf.sin(x) * tf.sin(y) + tf.cos(x) * (a + b * tf.sin(y)), d * tf.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return tf.convert_to_tensor( - [ - [ - -c * tf.cos(x) * tf.sin(y) - tf.sin(x) * (a + b * tf.sin(y)), - b * tf.cos(x) * tf.cos(y) - c * tf.cos(y) * tf.sin(x), - tf.cos(x), - tf.cos(x) * tf.sin(y), - -(tf.sin(x) * tf.sin(y)), - 0, - ], - [-d * tf.sin(x), 0, 0, 0, 0, tf.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - device_vjp = execute_kwargs.get("device_vjp", False) - - coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) - coeffs2 = tf.constant([0.7], dtype=tf.float64) - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost_fn(weights, coeffs1, coeffs2) - - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = tape.jacobian(res, [weights], experimental_use_pfor=not device_vjp) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) - coeffs2 = tf.Variable([0.7], dtype=tf.float64) - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - - with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2) - - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = qml.math.hstack(tape.jacobian(res, [weights, coeffs1, coeffs2]), like="tensorflow") - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) - - -@pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) -def test_device_returns_float32(diff_method): - """Test that if the device returns float32, the derivative succeeds.""" - - def _to_float32(results): - if isinstance(results, (list, tuple)): - return tuple(_to_float32(r) for r in results) - return np.array(results, dtype=np.float32) - - class Float32Dev(qml.devices.DefaultQubit): - def execute(self, circuits, execution_config=qml.devices.DefaultExecutionConfig): - results = super().execute(circuits, execution_config) - return _to_float32(results) - - dev = Float32Dev() - - @qml.qnode(dev, diff_method=diff_method) - def circuit(x): - qml.RX(tf.cos(x), wires=0) - return qml.expval(qml.Z(0)) - - x = tf.Variable(0.1, dtype=tf.float64) - - with tf.GradientTape() as tape: - y = circuit(x) - - assert qml.math.allclose(y, np.cos(np.cos(0.1))) - - g = tape.gradient(y, x) - expected_g = np.sin(np.cos(0.1)) * np.sin(0.1) - assert qml.math.allclose(g, expected_g) diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py deleted file mode 100644 index 3cdcf5eae30..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py +++ /dev/null @@ -1,849 +0,0 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Torch specific tests for execute and default qubit 2.""" -import numpy as np -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -torch = pytest.importorskip("torch") - -pytestmark = pytest.mark.torch - - -@pytest.fixture(autouse=True) -def run_before_and_after_tests(): - torch.set_default_dtype(torch.float64) - yield - torch.set_default_dtype(torch.float32) - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.skip("caching is not implemented for torch") - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = torch.arange(1, num_params + 1, requires_grad=True, dtype=torch.float64) - - N = len(params) - - def get_cost_tape(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - return qml.tape.QuantumScript.from_queue(q) - - def cost_no_cache(x): - return qml.execute( - [get_cost_tape(x)], - dev, - gradient_fn=qml.gradients.param_shift, - cache=False, - max_diff=2, - )[0] - - def cost_cache(x): - return qml.execute( - [get_cost_tape(x)], - dev, - gradient_fn=qml.gradients.param_shift, - cache=True, - max_diff=2, - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - hess1 = torch.autograd.functional.hessian(cost_no_cache, params) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params.clone().detach() - expected = torch.tensor( - [ - [2 * torch.cos(2 * x) * torch.sin(y) ** 2, torch.sin(2 * x) * torch.sin(2 * y)], - [ - torch.sin(2 * x) * torch.sin(2 * y), - -2 * torch.cos(x) ** 2 * torch.cos(2 * y), - ], - ] - ) - assert torch.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - hess2 = torch.autograd.functional.hessian(cost_cache, params) - assert torch.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), - ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, - Shots(None), - DefaultQubit(), - ), - ( - { - "gradient_fn": "adjoint", - "grad_on_execution": False, - "device_vjp": False, - }, - Shots(None), - DefaultQubit(), - ), - ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "device", "device_vjp": False}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(), - ), - ( - {"gradient_fn": "device", "device_vjp": True}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(), - ), -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestTorchExecuteIntegration: - """Test the torch interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=False) - - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("grad_on_execution", False): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - exp = torch.cos(a) * torch.cos(b) - if shots.has_partitioned_shots: - for shot in range(2): - for wire in range(2): - assert qml.math.allclose(res[shot][wire], exp, atol=atol_for_shots(shots)) - else: - for wire in range(2): - assert qml.math.allclose(res[wire], exp, atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = torch.tensor(0.1, requires_grad=True) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = torch.autograd.functional.jacobian(cost, a) - if not shots.has_partitioned_shots: - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - if shots.has_partitioned_shots: - for i in range(shots.num_copies): - assert torch.allclose(res[i], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[i], -torch.sin(a), atol=atol_for_shots(shots)) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res, -torch.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=True) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - [res] = execute([tape], device, **execute_kwargs) - if shots.has_partitioned_shots: - return torch.hstack(res[0] + res[1]) - return torch.hstack(res) - - res = cost(a, b) - expected = torch.tensor([torch.cos(a), -torch.cos(a) * torch.sin(b)]) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (a, b)) - assert isinstance(res, tuple) and len(res) == 2 - - expected = ( - torch.tensor([-torch.sin(a), torch.sin(a) * torch.sin(b)]), - torch.tensor([0, -torch.cos(a) * torch.cos(b)]), - ) - if shots.has_partitioned_shots: - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - for _r, _e in zip(res, expected): - assert torch.allclose(_r[:2], _e, atol=atol_for_shots(shots)) - assert torch.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - - else: - assert res[0].shape == (2,) - assert res[1].shape == (2,) - - for _r, _e in zip(res, expected): - assert torch.allclose(_r, _e, atol=atol_for_shots(shots)) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = [qml.math.asarray(ri, like="torch") for r in res for ri in r] - else: - res = [qml.math.asarray(r, like="torch") for r in res] - return sum(torch.hstack(res)) - - params = torch.tensor([0.1, 0.2], requires_grad=True) - x, y = params.clone().detach() - - res = cost(params) - expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - - if shots.has_partitioned_shots: - assert torch.allclose(res, 2 * expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res.backward() - expected = torch.tensor([-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)]) - if shots.has_partitioned_shots: - assert torch.allclose(params.grad, 2 * expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(params.grad, expected, atol=atol_for_shots(shots), rtol=0) - - @pytest.mark.skip("torch cannot reuse tensors in various computations") - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if execute_kwargs["gradient_fn"] == "backprop": - pytest.xfail("backprop is not compatible with something about this situation.") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = execute([tape1, tape2, tape3], device, **execute_kwargs) - return torch.hstack([qml.math.asarray(r, like="torch") for r in res]) - - params = torch.tensor([0.1, 0.2], requires_grad=True) - x, y = params.clone().detach() - - res = cost(params) - assert isinstance(res, torch.Tensor) - assert res.shape == (4,) - - assert torch.allclose(res[0], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - assert torch.allclose(res[1], torch.tensor(1.0), atol=atol_for_shots(shots)) - assert torch.allclose(res[2], torch.cos(torch.tensor(0.5)), atol=atol_for_shots(shots)) - assert torch.allclose(res[3], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - - jac = torch.autograd.functional.jacobian(cost, params) - assert isinstance(jac, torch.Tensor) - assert jac.shape == (4, 2) # pylint: disable=no-member - - assert torch.allclose(jac[1:3], torch.tensor(0.0), atol=atol_for_shots(shots)) - - d1 = -torch.sin(x) * torch.cos(y) - assert torch.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) # fails for torch - assert torch.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -torch.cos(x) * torch.sin(y) - assert torch.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) # fails for torch - assert torch.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=True) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return torch.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac = torch.autograd.functional.jacobian(cost, (a, b)) - - a = torch.tensor(0.54, requires_grad=True) - b = torch.tensor(0.8, requires_grad=True) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = torch.tensor([torch.cos(2 * a), -torch.cos(2 * a) * torch.sin(b)]) - assert torch.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac = torch.autograd.functional.jacobian(lambda a, b: cost(2 * a, b), (a, b)) - expected = ( - torch.tensor([-2 * torch.sin(2 * a), 2 * torch.sin(2 * a) * torch.sin(b)]), - torch.tensor([0, -torch.cos(2 * a) * torch.cos(b)]), - ) - assert isinstance(jac, tuple) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert torch.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=False) - c = torch.tensor(0.3, requires_grad=True) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + torch.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - # PyTorch docs suggest a lambda for cost functions with some non-trainable args - # See for more: https://pytorch.org/docs/stable/autograd.html#functional-higher-level-api - res = torch.autograd.functional.jacobian(lambda _a, _c: cost(_a, b, _c), (a, c)) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - @pytest.mark.skip("torch handles gradients and jacobians differently") - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = torch.tensor(0.1, requires_grad=False) - b = torch.tensor(0.2, requires_grad=False) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - assert res.shape == (2,) - - res = torch.autograd.functional.jacobian(cost, (a, b)) - assert len(res) == 0 - - def loss(a, b): - return torch.sum(cost(a, b)) - - res = loss(a, b) - res.backward() - - assert torch.allclose(torch.tensor([a.grad, b.grad]), 0) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the torch interface works correctly - with a matrix parameter""" - U = torch.tensor([[0, 1], [1, 0]], requires_grad=False, dtype=torch.float64) - a = torch.tensor(0.1, requires_grad=True) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert torch.allclose(res, -torch.cos(a), atol=atol_for_shots(shots)) - - jac = torch.autograd.functional.jacobian(lambda y: cost(y, U), a) - assert isinstance(jac, torch.Tensor) - assert torch.allclose(jac, torch.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = torch.tensor(0.1, requires_grad=False) - p = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - - res = cost_fn(a, p) - expected = torch.cos(a) * torch.cos(p[1]) * torch.sin(p[0]) + torch.sin(a) * ( - torch.cos(p[2]) * torch.sin(p[1]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.sin(p[2]) - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(lambda _p: cost_fn(a, _p), p) - expected = torch.tensor( - [ - torch.cos(p[1]) - * ( - torch.cos(a) * torch.cos(p[0]) - - torch.sin(a) * torch.sin(p[0]) * torch.sin(p[2]) - ), - torch.cos(p[1]) * torch.cos(p[2]) * torch.sin(a) - - torch.sin(p[1]) - * ( - torch.cos(a) * torch.sin(p[0]) - + torch.cos(p[0]) * torch.sin(a) * torch.sin(p[2]) - ), - torch.sin(a) - * ( - torch.cos(p[0]) * torch.cos(p[1]) * torch.cos(p[2]) - - torch.sin(p[1]) * torch.sin(p[2]) - ), - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = torch.tensor(0.543, requires_grad=True) - y = torch.tensor(-0.654, requires_grad=True) - - res = cost(x, y) - expected = torch.tensor( - [ - [ - torch.cos(x / 2) ** 2, - torch.sin(x / 2) ** 2, - (1 + torch.cos(x) * torch.cos(y)) / 2, - (1 - torch.cos(x) * torch.cos(y)) / 2, - ], - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (x, y)) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - torch.tensor( - [ - [ - -torch.sin(x) / 2, - torch.sin(x) / 2, - -torch.sin(x) * torch.cos(y) / 2, - torch.sin(x) * torch.cos(y) / 2, - ], - ] - ), - torch.tensor( - [ - [0, 0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2], - ] - ), - ) - - assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = torch.tensor(0.543, requires_grad=True) - y = torch.tensor(-0.654, requires_grad=True) - - res = cost(x, y) - expected = torch.tensor( - [ - torch.cos(x), - (1 + torch.cos(x) * torch.cos(y)) / 2, - (1 - torch.cos(x) * torch.cos(y)) / 2, - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (x, y)) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - torch.tensor( - [-torch.sin(x), -torch.sin(x) * torch.cos(y) / 2, torch.sin(x) * torch.cos(y) / 2] - ), - torch.tensor([0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2]), - ) - assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the torch execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64), - torch.tensor([0, -0.654], requires_grad=True, dtype=torch.float64), - torch.tensor([-2.0, 0], requires_grad=True, dtype=torch.float64), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using torch, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params.clone().detach() - expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - res.backward() - expected = torch.tensor( - [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] - ) - assert torch.allclose(params.grad, expected, atol=tol, rtol=0) - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.tensor( - [ - [-torch.cos(2 * x) * torch.cos(2 * y), torch.sin(2 * x) * torch.sin(2 * y)], - [torch.sin(2 * x) * torch.sin(2 * y), -2 * torch.cos(x) ** 2 * torch.cos(2 * y)], - ] - ) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = torch.tensor([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params.clone().detach() - expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - res.backward() - expected = torch.tensor( - [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] - ) - assert torch.allclose(params.grad, expected, atol=tol, rtol=0) - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2]) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - res = execute([tape], device, **execute_kwargs)[0] - if shots.has_partitioned_shots: - return torch.hstack(res[0] + res[1]) - return torch.hstack(res) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1.clone().detach() - d = coeffs2[0].clone().detach() - x, y = weights.clone().detach() - return torch.tensor( - [ - -c * torch.sin(x) * torch.sin(y) + torch.cos(x) * (a + b * torch.sin(y)), - d * torch.cos(x), - ] - ) - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.clone().detach() - d = coeffs2[0].clone().detach() - x, y = weights.clone().detach() - return torch.tensor( - [ - [ - -c * torch.cos(x) * torch.sin(y) - torch.sin(x) * (a + b * torch.sin(y)), - b * torch.cos(x) * torch.cos(y) - c * torch.cos(y) * torch.sin(x), - torch.cos(x), - torch.cos(x) * torch.sin(y), - -(torch.sin(x) * torch.sin(y)), - 0, - ], - [-d * torch.sin(x), 0, 0, 0, 0, torch.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False) - coeffs2 = torch.tensor([0.7], requires_grad=False) - weights = torch.tensor([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(lambda w: cost_fn(w, coeffs1, coeffs2), weights) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert torch.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - coeffs2 = torch.tensor([0.7], requires_grad=True) - weights = torch.tensor([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.hstack(torch.autograd.functional.jacobian(cost_fn, (weights, coeffs1, coeffs2))) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py new file mode 100644 index 00000000000..311c2f32b21 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py @@ -0,0 +1,1552 @@ +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the autograd interface""" +# pylint: disable=protected-access,too-few-public-methods +import sys + +import autograd +import pytest + +import pennylane as qml +from pennylane import numpy as np +from pennylane.devices import DefaultQubitLegacy +from pennylane.gradients import finite_diff, param_shift +from pennylane.operation import AnyWires, Observable + +pytestmark = pytest.mark.autograd + + +class TestAutogradExecuteUnitTests: + """Unit tests for autograd execution""" + + def test_import_error(self, mocker): + """Test that an exception is caught on import error""" + + mock = mocker.patch.object(autograd.extend, "defvjp") + mock.side_effect = ImportError() + + try: + del sys.modules["pennylane.workflow.interfaces.autograd"] + except KeyError: + pass + + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + + with qml.queuing.AnnotatedQueue() as q: + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.raises( + qml.QuantumFunctionError, + match="autograd not found. Please install the latest version " + "of autograd to enable the 'autograd' interface", + ): + qml.execute([tape], dev, gradient_fn=param_shift, interface="autograd") + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + qml.jacobian(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if a gradient transform + is used with grad_on_execution=True""" + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, gradient_fn=param_shift, grad_on_execution=True)[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + qml.jacobian(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, gradient_fn=param_shift, interface="None")[0] + + with pytest.raises(ValueError, match="interface must be in"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + + a = np.array([0.1, 0.2], requires_grad=True) + cost(a) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_no_gradients_on_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = np.array([0.1, 0.2], requires_grad=True) + cost(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + qml.jacobian(cost)(a) + spy_gradients.assert_called() + + +class TestBatchTransformExecution: + """Tests to ensure batch transforms can be correctly executed + via qml.execute and batch_transform""" + + def test_no_batch_transform(self, mocker): + """Test that batch transforms can be disabled and enabled""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100000) + + H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0) + x = 0.6 + y = 0.2 + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=0) + qml.RY(y, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.spy(dev, "batch_transform") + + if not qml.operation.active_new_opmath(): + with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"): + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + _ = qml.execute([tape], dev, None, device_batch_transform=False) + else: + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + res = qml.execute([tape], dev, None, device_batch_transform=False) + assert np.allclose(res[0], np.cos(y), atol=0.1) + + spy.assert_not_called() + + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + res = qml.execute([tape], dev, None, device_batch_transform=True) + spy.assert_called() + + assert qml.math.shape(res[0]) == () + assert np.allclose(res[0], np.cos(y), rtol=0.05) + + def test_batch_transform_dynamic_shots(self): + """Tests that the batch transform considers the number of shots for the execution, not those + statically on the device.""" + dev = qml.device("default.qubit.legacy", wires=1) + H = 2.0 * qml.PauliZ(0) + qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)]) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([qscript], dev, interface=None, override_shots=10) + assert res == (2.0,) + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cachesize=cachesize)[0] + + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] + + custom_cache = {} + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] + + # Without caching, 5 jacobians should still be performed + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = qml.jacobian(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.arange(1, num_params + 1) / 10 + + N = len(params) + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache, max_diff=2)[0] + + # No caching: number of executions is not ideal + hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params + expected = np.array( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert dev.num_executions == expected_runs + + # Use caching: number of executions is ideal + dev._num_executions = 0 + hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + def test_caching_adjoint_no_grad_on_execution(self): + """Test that caching reduces the number of adjoint evaluations + when the grads is not on execution.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack( + qml.execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + ) + + # no caching, but jac for each batch still stored. + qml.jacobian(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, only 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = qml.jacobian(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + def test_single_backward_pass_batch(self): + """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift + is requested for multiple tapes.""" + + dev = qml.device("default.qubit.legacy", wires=2) + + def f(x): + tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) + tape2 = qml.tape.QuantumScript([qml.RY(x, 0)], [qml.probs(wires=0)]) + + results = qml.execute([tape1, tape2], dev, gradient_fn=qml.gradients.param_shift) + return results[0] + results[1] + + x = qml.numpy.array(0.1) + with dev.tracker: + out = qml.jacobian(f)(x) + + assert dev.tracker.totals["batches"] == 2 + assert dev.tracker.history["batch_len"] == [2, 4] + expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] + assert qml.math.allclose(out, expected) + + def test_single_backward_pass_split_hamiltonian(self): + """Tests that the backward pass is one single batch, not a bunch of batches, when parameter + shift derivatives are requested for a tape that the device split into batches.""" + + dev = qml.device("default.qubit.legacy", wires=2, shots=50000) + + H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc") + + def f(x): + tape = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.expval(H)]) + return qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift)[0] + + x = qml.numpy.array(0.1) + with dev.tracker: + out = qml.grad(f)(x) + + assert dev.tracker.totals["batches"] == 2 + assert dev.tracker.history["batch_len"] == [2, 4] + + assert qml.math.allclose(out, -np.cos(x) - np.sin(x), atol=0.05) + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestAutogradExecuteIntegration: + """Test the autograd interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], dev, **execute_kwargs) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=False) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = np.array(0.1, requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, **execute_kwargs)[0] + + res = qml.jacobian(cost)(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_jacobian(self, execute_kwargs, tol): + """Test jacobian calculation""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost(a, b, device=dev) + expected = [np.cos(a), -np.cos(a) * np.sin(b)] + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(cost)(a, b, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (2,) + assert res[1].shape == (2,) + + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected)) + + @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") + def test_tape_no_parameters(self, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(params): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + with qml.queuing.AnnotatedQueue() as q4: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.probs(wires=[0, 1]) + + tape4 = qml.tape.QuantumScript.from_queue(q4) + return sum( + autograd.numpy.hstack( + qml.execute([tape1, tape2, tape3, tape4], dev, **execute_kwargs) + ) + ) + + params = np.array([0.1, 0.2], requires_grad=True) + x, y = params + + res = cost(params) + expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = qml.grad(cost)(params) + expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] + assert np.allclose(grad, expected, atol=tol, rtol=0) + + @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") + def test_tapes_with_different_return_size(self, execute_kwargs): + """Test that tapes wit different can be executed and differentiated.""" + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(params): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + return autograd.numpy.hstack(qml.execute([tape1, tape2, tape3], dev, **execute_kwargs)) + + params = np.array([0.1, 0.2], requires_grad=True) + + res = cost(params) + assert isinstance(res, np.ndarray) + assert res.shape == (4,) + + jac = qml.jacobian(cost)(params) + assert isinstance(jac, np.ndarray) + assert jac.shape == (4, 2) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return autograd.numpy.hstack(qml.execute([new_tape], dev, **execute_kwargs)[0]) + + jac_fn = qml.jacobian(cost) + jac = jac_fn(a, b) + + a = np.array(0.54, requires_grad=True) + b = np.array(0.8, requires_grad=True) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) + jac = jac_fn(a, b) + expected = ( + [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)], + [0, -np.cos(2 * a) * np.cos(b)], + ) + assert isinstance(jac, tuple) and len(jac) == 2 + assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected)) + + def test_classical_processing(self, execute_kwargs): + """Test classical processing within the quantum tape""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=False) + c = np.array(0.3, requires_grad=True) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + np.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = qml.jacobian(cost)(a, b, c, device=dev) + + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_no_trainable_parameters(self, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + a = np.array(0.1, requires_grad=False) + b = np.array(0.2, requires_grad=False) + + def cost(a, b, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, b, device=dev) + assert res.shape == (2,) + + with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): + res = qml.jacobian(cost)(a, b, device=dev) + assert len(res) == 0 + + def loss(a, b): + return np.sum(cost(a, b, device=dev)) + + with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): + res = qml.grad(loss)(a, b) + + assert np.allclose(res, 0) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the autograd interface works correctly + with a matrix parameter""" + U = np.array([[0, 1], [1, 0]], requires_grad=False) + a = np.array(0.1, requires_grad=True) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, U, device=dev) + assert isinstance(res, np.ndarray) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + jac = jac_fn(a, U, device=dev) + assert isinstance(jac, np.ndarray) + assert np.allclose(jac, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + with qml.queuing.AnnotatedQueue() as q_tape: + qml.RX(a, wires=0) + U3(*p, wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q_tape) + tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return qml.execute([tape], device, **execute_kwargs)[0] + + a = np.array(0.1, requires_grad=False) + p = np.array([0.1, 0.2, 0.3], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + res = cost_fn(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost_fn) + res = jac_fn(a, p, device=dev) + expected = np.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + def cost(x, y, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + x = np.array(0.543, requires_grad=True) + y = np.array(-0.654, requires_grad=True) + + res = cost(x, y, device=dev) + expected = np.array( + [ + [ + np.cos(x / 2) ** 2, + np.sin(x / 2) ** 2, + (1 + np.cos(x) * np.cos(y)) / 2, + (1 - np.cos(x) * np.cos(y)) / 2, + ], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + res = jac_fn(x, y, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ( + np.array( + [ + [ + -np.sin(x) / 2, + np.sin(x) / 2, + -np.sin(x) * np.cos(y) / 2, + np.sin(x) * np.cos(y) / 2, + ], + ] + ), + np.array( + [ + [0, 0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], + ] + ), + ) + + assert np.allclose(res[0], expected[0], atol=tol, rtol=0) + assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + def cost(x, y, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + x = np.array(0.543, requires_grad=True) + y = np.array(-0.654, requires_grad=True) + + res = cost(x, y, device=dev) + expected = np.array( + [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + res = jac_fn(x, y, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = ( + np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), + np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), + ) + assert np.allclose(res[0], expected[0], atol=tol, rtol=0) + assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + + def test_sampling(self, execute_kwargs): + """Test sampling works as expected""" + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + def cost(device): + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) + res = cost(device=dev) + assert isinstance(res, tuple) + assert len(res) == 2 + assert res[0].shape == (shots,) + assert res[1].shape == (shots,) + + +class TestHigherOrderDerivatives: + """Test that the autograd execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + np.array([0.543, -0.654], requires_grad=True), + np.array([0, -0.654], requires_grad=True), + np.array([-2.0, 0], requires_grad=True), + ], + ) + def test_parameter_shift_hessian(self, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using autograd, yielding second derivatives.""" + dev = qml.device("default.qubit.autograd", wires=2) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.grad(cost_fn)(params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(qml.grad(cost_fn))(params) + expected = np.array( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_adjoint_hessian_one_param(self, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.autograd", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.grad(cost_fn))(params) + + assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0) + + def test_adjoint_hessian_multiple_params(self, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.autograd", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack( + qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + ) + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.jacobian(cost_fn))(params) + + assert np.allclose(res, np.zeros([2, 2, 2]), atol=tol, rtol=0) + + def test_max_diff(self, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.grad(cost_fn)(params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.grad(cost_fn))(params) + + expected = np.zeros([2, 2]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +class TestOverridingShots: + """Test overriding shots on execution""" + + def test_changing_shots(self, mocker, tol): + """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.spy(dev, "sample") + + # execute with device default shots (None) + res = qml.execute([tape], dev, gradient_fn=param_shift) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() + + # execute with shots=100 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = qml.execute([tape], dev, gradient_fn=param_shift) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # same single call from above, no additional calls + + def test_overriding_shots_with_same_value(self, mocker): + """Overriding shots with the same value as the device will have no effect""" + dev = qml.device("default.qubit.legacy", wires=2, shots=123) + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.Mock(wraps=qml.Device.shots.fset) + # pylint:disable=assignment-from-no-return,too-many-function-args + mock_property = qml.Device.shots.setter(spy) + mocker.patch.object(qml.Device, "shots", mock_property) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([tape], dev, gradient_fn=param_shift, override_shots=123) + # overriden shots is the same, no change + spy.assert_not_called() + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) + # overriden shots is not the same, shots were changed + spy.assert_called() + + # shots were temporarily set to the overriden value + assert spy.call_args_list[0][0] == (dev, 100) + # shots were then returned to the built-in value + assert spy.call_args_list[1][0] == (dev, 123) + + def test_overriding_device_with_shot_vector(self): + """Overriding a device that has a batch of shots set + results in original shots being returned after execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=[10, (1, 3), 5]) + + assert dev.shots == 18 + assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] + + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100)[0] + + assert isinstance(res, np.ndarray) + assert res.shape == () + + # device is unchanged + assert dev.shots == 18 + assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] + + res = qml.execute([tape], dev, gradient_fn=param_shift)[0] + assert len(res) == 5 + + @pytest.mark.xfail(reason="Shots vector must be adapted for new return types.") + def test_gradient_integration(self): + """Test that temporarily setting the shots works + for gradient computations""" + # TODO: Update here when shot vectors are supported + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + a, b = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(a, b, shots): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + result = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=shots) + return result[0] + + res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000]) + assert dev.shots is None + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] + assert all( + np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected) + ) + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift}, + {"gradient_fn": finite_diff}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + """Test hamiltonian with no trainable parameters.""" + # pylint: disable=unused-argument + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) + coeffs2 = np.array([0.7], requires_grad=False) + weights = np.array([0.4, 0.5], requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + """Test hamiltonian with trainable parameters.""" + # pylint: disable=unused-argument + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) + coeffs2 = np.array([0.7], requires_grad=True) + weights = np.array([0.4, 0.5], requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +class TestCustomJacobian: + """Test for custom Jacobian.""" + + def test_custom_jacobians(self): + """Test custom Jacobian device methood""" + + class CustomJacobianDevice(DefaultQubitLegacy): + @classmethod + def capabilities(cls): + capabilities = super().capabilities() + capabilities["provides_jacobian"] = True + return capabilities + + def jacobian(self, tape): + # pylint: disable=unused-argument + return np.array([1.0, 2.0, 3.0, 4.0]) + + dev = CustomJacobianDevice(wires=2) + + @qml.qnode(dev, diff_method="device") + def circuit(v): + qml.RX(v, wires=0) + return qml.probs(wires=[0, 1]) + + d_circuit = qml.jacobian(circuit, argnum=0) + + params = np.array(1.0, requires_grad=True) + + d_out = d_circuit(params) + assert np.allclose(d_out, np.array([1.0, 2.0, 3.0, 4.0])) + + def test_custom_jacobians_param_shift(self): + """Test computing the gradient using the parameter-shift + rule with a device that provides a jacobian""" + + class MyQubit(DefaultQubitLegacy): + @classmethod + def capabilities(cls): + capabilities = super().capabilities().copy() + capabilities.update( + provides_jacobian=True, + ) + return capabilities + + def jacobian(self, *args, **kwargs): + raise NotImplementedError() + + dev = MyQubit(wires=2) + + @qml.qnode(dev, diff_method="parameter-shift", grad_on_execution=False) + def qnode(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b): + return autograd.numpy.hstack(qnode(a, b)) + + res = qml.jacobian(cost)(a, b) + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + + assert np.allclose(res[0], expected[0]) + assert np.allclose(res[1], expected[1]) + + +class SpecialObject: + """SpecialObject + A special object that conveniently encapsulates the return value of + a special observable supported by a special device and which supports + multiplication with scalars and addition. + """ + + def __init__(self, val): + self.val = val + + def __mul__(self, other): + new = SpecialObject(self.val) + new *= other + return new + + def __imul__(self, other): + self.val *= other + return self + + def __rmul__(self, other): + return self * other + + def __iadd__(self, other): + self.val += other.val if isinstance(other, self.__class__) else other + return self + + def __add__(self, other): + new = SpecialObject(self.val) + new += other.val if isinstance(other, self.__class__) else other + return new + + def __radd__(self, other): + return self + other + + +class SpecialObservable(Observable): + """SpecialObservable""" + + num_wires = AnyWires + num_params = 0 + par_domain = None + + def diagonalizing_gates(self): + """Diagonalizing gates""" + return [] + + +class DeviceSupportingSpecialObservable(DefaultQubitLegacy): + name = "Device supporting SpecialObservable" + short_name = "default.qubit.specialobservable" + observables = DefaultQubitLegacy.observables.union({"SpecialObservable"}) + + @staticmethod + def _asarray(arr, dtype=None): + # pylint: disable=unused-argument + return arr + + @classmethod + def capabilities(cls): + capabilities = super().capabilities().copy() + capabilities.update( + provides_jacobian=True, + ) + return capabilities + + def expval(self, observable, **kwargs): + if self.analytic and isinstance(observable, SpecialObservable): + val = super().expval(qml.PauliZ(wires=0), **kwargs) + return np.array(SpecialObject(val)) + + return super().expval(observable, **kwargs) + + def jacobian(self, tape): + # we actually let pennylane do the work of computing the + # jacobian for us but return it as a device jacobian + gradient_tapes, fn = qml.gradients.param_shift(tape) + tape_jacobian = fn(qml.execute(gradient_tapes, self, None)) + return tape_jacobian + + +@pytest.mark.autograd +class TestObservableWithObjectReturnType: + """Unit tests for qnode returning a custom object""" + + def test_custom_return_type(self): + """Test custom return values for a qnode""" + + dev = DeviceSupportingSpecialObservable(wires=1, shots=None) + + # force diff_method='parameter-shift' because otherwise + # PennyLane swaps out dev for default.qubit.autograd + @qml.qnode(dev, diff_method="parameter-shift") + def qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + @qml.qnode(dev, diff_method="parameter-shift") + def reference_qnode(x): + qml.RY(x, wires=0) + return qml.expval(qml.PauliZ(wires=0)) + + out = qnode(0.2) + assert isinstance(out, np.ndarray) + assert isinstance(out.item(), SpecialObject) + assert np.isclose(out.item().val, reference_qnode(0.2)) + + def test_jacobian_with_custom_return_type(self): + """Test differentiation of a QNode on a device supporting a + special observable that returns an object rather than a number.""" + + dev = DeviceSupportingSpecialObservable(wires=1, shots=None) + + # force diff_method='parameter-shift' because otherwise + # PennyLane swaps out dev for default.qubit.autograd + @qml.qnode(dev, diff_method="parameter-shift") + def qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + @qml.qnode(dev, diff_method="parameter-shift") + def reference_qnode(x): + qml.RY(x, wires=0) + return qml.expval(qml.PauliZ(wires=0)) + + reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),) + + assert np.isclose( + reference_jac, + qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val, + ) + + # now check that also the device jacobian works with a custom return type + @qml.qnode(dev, diff_method="device") + def device_gradient_qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + assert np.isclose( + reference_jac, + qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val, + ) diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py similarity index 69% rename from tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py index a0d1e0f2e66..818d7d58ff5 100644 --- a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2020 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the autograd interface with a QNode""" -# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-dict-literal, no-name-in-module - +# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison, +# pylint: disable=unnecessary-lambda-assignment import autograd import autograd.numpy as anp import pytest @@ -21,47 +21,32 @@ import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit -from tests.param_shift_dev import ParamShiftDerivativesDevice -# dev, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [qml.device("default.qubit"), "finite-diff", False, False], - [qml.device("default.qubit"), "parameter-shift", False, False], - [qml.device("default.qubit"), "backprop", True, False], - [qml.device("default.qubit"), "adjoint", True, False], - [qml.device("default.qubit"), "adjoint", False, False], - [qml.device("default.qubit"), "spsa", False, False], - [qml.device("default.qubit"), "hadamard", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, False], - [ParamShiftDerivativesDevice(), "best", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_qubit_device_and_diff_method = [ - ["autograd", DefaultQubit(), "finite-diff", False, False], - ["autograd", DefaultQubit(), "parameter-shift", False, False], - ["autograd", DefaultQubit(), "backprop", True, False], - ["autograd", DefaultQubit(), "adjoint", True, False], - ["autograd", DefaultQubit(), "adjoint", False, False], - ["autograd", DefaultQubit(), "adjoint", True, True], - ["autograd", DefaultQubit(), "adjoint", False, True], - ["autograd", DefaultQubit(), "spsa", False, False], - ["autograd", DefaultQubit(), "hadamard", False, False], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, True], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, True], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, False], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, False], - ["auto", DefaultQubit(), "finite-diff", False, False], - ["auto", DefaultQubit(), "parameter-shift", False, False], - ["auto", DefaultQubit(), "backprop", True, False], - ["auto", DefaultQubit(), "adjoint", True, False], - ["auto", DefaultQubit(), "adjoint", False, False], - ["auto", DefaultQubit(), "spsa", False, False], - ["auto", DefaultQubit(), "hadamard", False, False], - ["auto", qml.device("lightning.qubit", wires=5), "adjoint", False, False], - ["auto", qml.device("lightning.qubit", wires=5), "adjoint", True, False], + ["autograd", "default.qubit.legacy", "finite-diff", False], + ["autograd", "default.qubit.legacy", "parameter-shift", False], + ["autograd", "default.qubit.legacy", "backprop", True], + ["autograd", "default.qubit.legacy", "adjoint", True], + ["autograd", "default.qubit.legacy", "adjoint", False], + ["autograd", "default.qubit.legacy", "spsa", False], + ["autograd", "default.qubit.legacy", "hadamard", False], + ["auto", "default.qubit.legacy", "finite-diff", False], + ["auto", "default.qubit.legacy", "parameter-shift", False], + ["auto", "default.qubit.legacy", "backprop", True], + ["auto", "default.qubit.legacy", "adjoint", True], + ["auto", "default.qubit.legacy", "adjoint", False], + ["auto", "default.qubit.legacy", "spsa", False], + ["auto", "default.qubit.legacy", "hadamard", False], ] pytestmark = pytest.mark.autograd @@ -72,21 +57,78 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with Autograd integrates with the PennyLane stack""" - # pylint:disable=unused-argument - def test_execution_no_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp + # pylint: disable=unused-argument + + def test_nondiff_param_unwrapping( + self, interface, dev_name, diff_method, grad_on_execution, mocker ): + """Test that non-differentiable parameters are correctly unwrapped + to NumPy ndarrays or floats (if 0-dimensional)""" + if diff_method != "parameter-shift": + pytest.skip("Test only supports parameter-shift") + + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method=diff_method) + def circuit(x, y): + qml.RX(x[0], wires=0) + qml.Rot(*x[1:], wires=0) + qml.RY(y[0], wires=0) + return qml.expval(qml.PauliZ(0)) + + x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False) + y = np.array([0.5], requires_grad=True) + + param_data = [] + + def mock_apply(*args, **kwargs): + for op in args[0]: + param_data.extend(op.data) + + mocker.patch.object(dev, "apply", side_effect=mock_apply) + circuit(x, y) + assert param_data == [0.1, 0.2, 0.3, 0.4, 0.5] + assert not any(isinstance(p, np.tensor) for p in param_data) + + # test the jacobian works correctly + param_data = [] + qml.grad(circuit)(x, y) + assert param_data == [ + 0.1, + 0.2, + 0.3, + 0.4, + 0.5, + 0.1, + 0.2, + 0.3, + 0.4, + 0.5 + np.pi / 2, + 0.1, + 0.2, + 0.3, + 0.4, + 0.5 - np.pi / 2, + ] + assert not any(isinstance(p, np.tensor) for p in param_data) + + def test_execution_no_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works without an interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - @qnode(dev, interface=None) + num_wires = 1 + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface=None, diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -96,24 +138,26 @@ def circuit(a): res = circuit(a) - # without the interface, the QNode simply returns a scalar array or float - assert isinstance(res, (np.ndarray, float)) - assert qml.math.shape(res) == tuple() # pylint: disable=comparison-with-callable + # without the interface, the QNode simply returns a scalar array + assert isinstance(res, np.ndarray) + assert res.shape == tuple() - def test_execution_with_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + # gradients should cause an error + with pytest.raises(TypeError, match="must be real number, not ArrayBox"): + qml.grad(circuit)(a) + + def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -121,27 +165,31 @@ def circuit(a): a = np.array(0.1, requires_grad=True) assert circuit.interface == interface + # gradients should work grad = qml.grad(circuit)(a) + assert isinstance(grad, float) assert grad.shape == tuple() - def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol, device_vjp): + def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation""" + num_wires = 2 kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -172,23 +220,24 @@ def cost(x, y): assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, interface, dev, diff_method, grad_on_execution, tol, device_vjp - ): + def test_jacobian_no_evaluate(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" + num_wires = 2 kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -214,31 +263,36 @@ def cost(x, y): assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): """Test setting jacobian options""" - if diff_method != "finite-diff": - pytest.skip("Test only supports finite diff.") + wires = [0] + if diff_method in ["backprop", "adjoint"]: + pytest.skip("Test does not support backprop or adjoint method") + elif diff_method == "finite-diff": + kwargs = {"h": 1e-8, "approx_order": 2} + elif diff_method == "parameter-shift": + kwargs = {"shifts": [(0.1,), (0.2,)]} + elif diff_method == "hadamard": + wires = [0, "aux"] + kwargs = {"aux_wire": qml.wires.Wires("aux"), "device_wires": wires} + else: + kwargs = {} a = np.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface=interface, - h=1e-8, - order=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device("default.qubit.legacy", wires=wires) + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) + circuit(a) + qml.jacobian(circuit)(a) - def test_changing_trainability( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -247,13 +301,9 @@ def test_changing_trainability( a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - @qnode( - dev, - interface=interface, - diff_method="parameter-shift", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -290,18 +340,21 @@ def loss(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -319,17 +372,12 @@ def circuit(a, b, c): assert res[0].shape == () assert res[1].shape == () - def test_no_trainable_parameters( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): """Test evaluation and Jacobian if there are no trainable parameters""" + dev = qml.device(dev_name, wires=2) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -348,34 +396,35 @@ def circuit(a, b): assert len(res) == 2 assert isinstance(res, tuple) - def cost(x, y): + def cost0(x, y): return autograd.numpy.hstack(circuit(x, y)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - assert not qml.jacobian(cost)(a, b) + assert not qml.jacobian(cost0)(a, b) - def cost2(a, b): + def cost1(a, b): return np.sum(circuit(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - grad = qml.grad(cost2)(a, b) + grad = qml.grad(cost1)(a, b) assert grad == tuple() - def test_matrix_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -391,18 +440,13 @@ def circuit(U, a): assert np.allclose(res, np.sin(a), atol=tol, rtol=0) def test_gradient_non_differentiable_exception( - self, interface, dev, diff_method, grad_on_execution, device_vjp + self, interface, dev_name, diff_method, grad_on_execution ): """Test that an exception is raised if non-differentiable data is differentiated""" + dev = qml.device(dev_name, wires=2) - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(data1): qml.templates.AmplitudeEmbedding(data1, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @@ -413,27 +457,18 @@ def circuit(data1): with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"): grad_fn(data1) - def test_differentiable_expand( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that operation and nested tape expansion is differentiable""" kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=20) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - # pylint: disable=too-few-public-methods class U3(qml.U3): - """Custom U3.""" - def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -442,6 +477,7 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + dev = qml.device(dev_name, wires=2) a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) @@ -455,14 +491,14 @@ def circuit(a, p): expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) res = qml.grad(circuit)(a, p) - # assert isinstance(res, np.ndarray) - # assert len(res) == 3 + assert isinstance(res, np.ndarray) + assert len(res) == 3 expected = np.array( [ @@ -481,9 +517,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and remains differentiable.""" - def test_changing_shots(self): + def test_changing_shots(self, mocker, tol): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -491,21 +527,31 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - res = circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) # pylint: disable=comparison-with-callable + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # same single call performed above @pytest.mark.xfail(reason="Param shift and shot vectors.") def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -528,57 +574,62 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = np.array([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit()) + @qnode(dev) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - assert cost_fn.gradient_fn == "backprop" # gets restored to default - - cost_fn(a, b, shots=100) + cost_fn(a, b) # since we are using finite shots, parameter-shift will # be chosen - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used - cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" + # pylint: disable=unused-argument + def test_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with a single prob output""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -599,23 +650,25 @@ def circuit(x, y): assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) def test_multiple_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -664,24 +717,22 @@ def cost(x, y): ) assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - def test_ragged_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + num_wires = 2 + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -717,25 +768,22 @@ def cost(x, y): assert np.allclose(res[1], expected[1], atol=tol, rtol=0) def test_ragged_differentiation_variance( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -756,12 +804,12 @@ def circuit(x, y): assert isinstance(res, tuple) assert len(res) == 2 - # assert isinstance(res[0], np.ndarray) - # assert res[0].shape == () + assert isinstance(res[0], np.ndarray) + assert res[0].shape == () assert np.allclose(res[0], expected_var, atol=tol, rtol=0) - # assert isinstance(res[1], np.ndarray) - # assert res[1].shape == (2,) + assert isinstance(res[1], np.ndarray) + assert res[1].shape == (2,) assert np.allclose(res[1], expected_probs, atol=tol, rtol=0) def cost(x, y): @@ -778,38 +826,33 @@ def cost(x, y): assert isinstance(jac, tuple) assert len(jac) == 2 - # assert isinstance(jac[0], np.ndarray) - # assert jac[0].shape == (3,) + assert isinstance(jac[0], np.ndarray) + assert jac[0].shape == (3,) assert np.allclose(jac[0], expected[0], atol=tol, rtol=0) - # assert isinstance(jac[1], np.ndarray) - # assert jac[1].shape == (3,) + assert isinstance(jac[1], np.ndarray) + assert jac[1].shape == (3,) assert np.allclose(jac[1], expected[1], atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" + num_wires = 2 - # pylint: disable=too-few-public-methods - class Template(qml.templates.StronglyEntanglingLayers): - """Custom template.""" + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit2(data, weights): qml.templates.AngleEmbedding(data, wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -833,22 +876,19 @@ def cost(w1, w2): assert len(res) == 2 - def test_chained_gradient_value( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_chained_gradient_value(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the returned gradient value for two chained qubit QNodes is correct.""" - kwargs = dict( - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - + kwargs = dict(interface=interface, diff_method=diff_method) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - dev1 = qml.device("default.qubit") + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev1 = qml.device(dev_name, wires=num_wires) @qnode(dev1, **kwargs) def circuit1(a, b, c): @@ -859,11 +899,9 @@ def circuit1(a, b, c): qml.CNOT(wires=[1, 2]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2)) - dev2 = dev + dev2 = qml.device("default.qubit.legacy", wires=num_wires) - @qnode( - dev2, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution - ) + @qnode(dev2, interface=interface, diff_method=diff_method) def circuit2(data, weights): qml.RX(data[0], wires=0) qml.RX(data[1], wires=1) @@ -933,20 +971,19 @@ def cost(a, b, c, weights): # to the first parameter of circuit1. assert circuit1.qtape.trainable_params == [1, 2] - def test_second_derivative( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_second_derivative(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -976,17 +1013,18 @@ def circuit(x): assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1001,8 +1039,8 @@ def circuit(x): expected_res = np.cos(a) * np.cos(b) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected_res, atol=tol, rtol=0) grad_fn = qml.grad(circuit) @@ -1030,18 +1068,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_unused_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1070,20 +1109,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1099,7 +1137,7 @@ def circuit(x): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) # pylint: disable=comparison-with-callable + assert res.shape == (2,) assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1136,18 +1174,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1190,18 +1229,19 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(a, b): @@ -1215,7 +1255,7 @@ def circuit(a, b): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) # pylint: disable=comparison-with-callable + assert res.shape == (2,) assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1232,12 +1272,8 @@ def circuit(a, b): assert g[1].shape == (2,) assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0) - def jac_fn_a(*args): - return jac_fn(*args)[0] - - def jac_fn_b(*args): - return jac_fn(*args)[1] - + jac_fn_a = lambda *args: jac_fn(*args)[0] + jac_fn_b = lambda *args: jac_fn(*args)[1] hess_a = qml.jacobian(jac_fn_a)(a, b) hess_b = qml.jacobian(jac_fn_b)(a, b) assert isinstance(hess_a, tuple) and len(hess_a) == 2 @@ -1259,17 +1295,18 @@ def jac_fn_b(*args): assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0) assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=2) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1283,18 +1320,14 @@ def circuit(x): a, b = x - expected_res = [ - np.cos(a) * np.cos(b), - 0.5 + 0.5 * np.cos(a) * np.cos(b), - 0.5 - 0.5 * np.cos(a) * np.cos(b), - ] + cos_prod = np.cos(a) * np.cos(b) + expected_res = (cos_prod, [0.5 + 0.5 * cos_prod, 0.5 - 0.5 * cos_prod]) + res = circuit(x) + assert all(qml.math.allclose(r, e) for r, e in zip(res, expected_res)) def cost_fn(x): return autograd.numpy.hstack(circuit(x)) - res = cost_fn(x) - assert qml.math.allclose(res, expected_res) - jac_fn = qml.jacobian(cost_fn) hess = qml.jacobian(jac_fn)(x) @@ -1318,21 +1351,18 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning does not support state adjoint diff.") + dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1363,25 +1393,20 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the variance of a projector is correctly returned""" - if diff_method == "adjoint": - pytest.skip("adjoint supports either expvals or diagonal measurements.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) P = np.array(state, requires_grad=False) x, y = np.array([0.765, -0.654], requires_grad=True) @@ -1394,8 +1419,8 @@ def circuit(x, y): res = circuit(x, y) expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) jac = qml.jacobian(circuit)(x, y) @@ -1419,60 +1444,106 @@ def circuit(x, y): assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + pytest.skip("Hadamard gradient does not support variances.") + + r = np.array(0.543, requires_grad=True) + phi = np.array(-0.654, requires_grad=True) + + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = qml.jacobian(circuit)(r, phi) + expected = np.array( + [ + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ] ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + assert np.allclose(res, expected, atol=tol, rtol=0) - phi = np.array(1.23, requires_grad=True) - theta = np.array(4.56, requires_grad=True) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + elif diff_method == "hadamard": + pytest.skip("Hadamard gradient does not support variances.") - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + n = np.array(0.12, requires_grad=True) + a = np.array(0.765, requires_grad=True) - gradient = qml.grad(circuit)(phi, theta) - exp_theta_grad = qml.grad(expected_circuit)(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method -) + # circuit jacobians + res = qml.jacobian(circuit)(n, a) + expected = np.array([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +def test_adjoint_reuse_device_state(mocker): + """Tests that the autograd interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + qml.grad(circ, argnum=0)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Autograd interface""" @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp + self, dev_name, diff_method, grad_on_execution, max_diff ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" @@ -1481,22 +1552,20 @@ def test_gradient_expansion_trainable_only( if max_diff == 2 and diff_method == "hadamard": pytest.skip("Max diff > 1 not supported for Hadamard gradient.") - # pylint: disable=too-few-public-methods - class PhaseShift(qml.PhaseShift): - """dummy phase shift.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): return [qml.RY(3 * self.data[0], wires=self.wires)] - @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - device_vjp=device_vjp, - ) + @qnode(dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff) def circuit(x, y): qml.Hadamard(wires=0) PhaseShift(x, wires=0) @@ -1507,11 +1576,11 @@ def circuit(x, y): y = np.array(0.7, requires_grad=False) circuit(x, y) - _ = qml.grad(circuit)(x, y) + qml.grad(circuit)(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, tol, device_vjp + self, dev_name, diff_method, grad_on_execution, max_diff, tol ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" @@ -1519,16 +1588,15 @@ def test_hamiltonian_expansion_analytic( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, ) - if diff_method in ["adjoint", "hadamard"]: pytest.skip("The diff method requested does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1558,11 +1626,7 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if ( - diff_method in ("parameter-shift", "backprop") - and max_diff == 2 - and dev.name != "param_shift.qubit" - ): + if diff_method in ("parameter-shift", "backprop") and max_diff == 2: if diff_method == "backprop": with pytest.warns(UserWarning, match=r"Output seems independent of input."): grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) @@ -1580,26 +1644,28 @@ def circuit(data, weights, coeffs): @pytest.mark.slow @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), - "num_directions": 10, - } + gradient_kwargs = dict( + h=H_FOR_SPSA, + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), + num_directions=20, + ) tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1607,7 +1673,6 @@ def test_hamiltonian_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1623,15 +1688,13 @@ def circuit(data, weights, coeffs): c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - - grad = qml.grad(circuit)(d, w, c, shots=50000) + grad = qml.grad(circuit)(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1641,11 +1704,12 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": - grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c, shots=50000) + if diff_method == "parameter-shift" and max_diff == 2: + with pytest.warns(UserWarning, match=r"Output seems independent of input."): + grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) assert np.allclose(grad2_c, 0, atol=tol) - grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c, shots=50000) + grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c) expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ 0, -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), @@ -1659,38 +1723,36 @@ class TestSample: def test_backprop_error(self): """Test that sampling in backpropagation grad_on_execution raises an error""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2) @qnode(dev, diff_method="backprop") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - with pytest.raises( - qml.QuantumFunctionError, match="does not support backprop with requested" - ): - circuit(shots=10) + with pytest.raises(qml.QuantumFunctionError, match="only supported when shots=None"): + circuit(shots=10) # pylint: disable=unexpected-keyword-arg def test_sample_dimension(self): """Test that the sample function outputs samples of the right size""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=n_sample) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 - assert res[0].shape == (10,) # pylint: disable=comparison-with-callable + assert res[0].shape == (10,) assert isinstance(res[0], np.ndarray) - assert res[1].shape == (10,) # pylint: disable=comparison-with-callable + assert res[1].shape == (10,) assert isinstance(res[1], np.ndarray) def test_sample_combination(self): @@ -1698,7 +1760,7 @@ def test_sample_combination(self): n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) @qnode(dev, diff_method="parameter-shift") def circuit(): @@ -1706,29 +1768,29 @@ def circuit(): return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=n_sample) + result = circuit() assert isinstance(result, tuple) assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert isinstance(result[1], (float, np.ndarray)) - assert isinstance(result[2], (float, np.ndarray)) + assert isinstance(result[1], np.ndarray) + assert isinstance(result[2], np.ndarray) assert result[0].dtype == np.dtype("float") def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=n_sample) + result = circuit() assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) @@ -1738,44 +1800,44 @@ def test_multi_wire_sample_regular_shape(self): where a rectangular array is expected""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=n_sample) + result = circuit() # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) assert len(result) == 3 - assert result[0].shape == (10,) # pylint: disable=comparison-with-callable + assert result[0].shape == (10,) assert isinstance(result[0], np.ndarray) - assert result[1].shape == (10,) # pylint: disable=comparison-with-callable + assert result[1].shape == (10,) assert isinstance(result[1], np.ndarray) - assert result[2].shape == (10,) # pylint: disable=comparison-with-callable + assert result[2].shape == (10,) assert isinstance(result[2], np.ndarray) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution,device_vjp", qubit_device_and_diff_method -) +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, device_vjp): + # pylint: disable=unused-argument + + def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution): """For one measurement and one param, the gradient is a float.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1785,20 +1847,25 @@ def circuit(a): grad = qml.grad(circuit)(a) - assert isinstance(grad, np.tensor if diff_method == "backprop" else float) + import sys - def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp - ): + python_version = sys.version_info.minor + if diff_method == "backprop" and python_version > 7: + # Since numpy 1.23.0 + assert isinstance(grad, np.ndarray) + else: + assert isinstance(grad, float) + + def test_grad_single_measurement_multiple_param(self, dev_name, diff_method, grad_on_execution): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1815,17 +1882,17 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1840,18 +1907,21 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1865,18 +1935,21 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1896,17 +1969,20 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single array is returned.""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1919,17 +1995,20 @@ def circuit(a): assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a single params return an array.""" + num_wires = 2 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1946,17 +2025,21 @@ def cost(x): assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1980,17 +2063,21 @@ def cost(x, y): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2006,8 +2093,14 @@ def cost(x): assert isinstance(jac, np.ndarray) assert jac.shape == (5, 2) - def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): + def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2015,14 +2108,7 @@ def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_executio par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2043,24 +2129,22 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_expval_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of single measurement with a multiple params array return a single array.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2072,8 +2156,9 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2083,14 +2168,7 @@ def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2111,25 +2189,18 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_var_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2141,24 +2212,19 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_probs_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2183,23 +2249,18 @@ def cost(x, y): assert hess[1].shape == (6,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2212,12 +2273,11 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) # pylint: disable=no-member + assert hess.shape == (3, 2, 2) - def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2227,14 +2287,7 @@ def test_hessian_probs_var_multiple_params( par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2258,25 +2311,18 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (6,) - def test_hessian_var_multiple_param_array2( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_var_probs_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2289,14 +2335,15 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) # pylint: disable=no-member + assert hess.shape == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = DefaultQubit() + dev = qml.device(dev_name, wires=1) @qml.qnode(dev, interface="autograd") def circuit(): diff --git a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py new file mode 100644 index 00000000000..24d1d824ced --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py @@ -0,0 +1,658 @@ +# Copyright 2022 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration tests for using the Autograd interface with shot vectors and with a QNode""" +# pylint: disable=too-many-arguments,redefined-outer-name + +import pytest + +import pennylane as qml +from pennylane import numpy as np +from pennylane import qnode + +pytestmark = pytest.mark.autograd + +shots_and_num_copies = [(((5, 2), 1, 10), 4), ((1, 10, (5, 2)), 4)] +shots_and_num_copies_hess = [(((5, 1), 10), 2), ((10, (5, 1)), 2)] + + +kwargs = { + "finite-diff": {"h": 0.05}, + "parameter-shift": {}, + "spsa": {"h": 0.05, "num_directions": 20}, +} + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff"], + ["default.qubit.legacy", "parameter-shift"], + ["default.qubit.legacy", "spsa"], +] + +TOLS = { + "finite-diff": 0.3, + "parameter-shift": 1e-2, + "spsa": 0.3, +} + + +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnWithShotVectors: + """Class to test the shape of the Jacobian/Hessian with different return types and shot vectors.""" + + def test_jac_single_measurement_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array(0.1) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies,) + + def test_jac_single_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array(0.1) + b = np.array(0.2) + + def cost(a, b): + return qml.math.stack(circuit(a, b)) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies,) + + def test_jacobian_single_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array([0.1, 0.2]) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2) + + def test_jacobian_single_measurement_param_probs( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single array is returned with the correct + dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array(0.1) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 4) + + def test_jacobian_single_measurement_probs_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array(0.1) + b = np.array(0.2) + + def cost(a, b): + return qml.math.stack(circuit(a, b)) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 4) + + def test_jacobian_single_measurement_probs_multiple_param_single_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array([0.1, 0.2]) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 4, 2) + + def test_jacobian_expval_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=1, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 2) + + def test_jacobian_expval_expval_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + a = np.array([0.1, 0.2, 0.3]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2, 3) + + def test_jacobian_var_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 2) + + def test_jacobian_var_var_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + a = np.array([0.1, 0.2, 0.3]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2, 3) + + def test_jacobian_multiple_measurement_single_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 5) + + def test_jacobian_multiple_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b): + res = circuit(a, b) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 5) + + def test_jacobian_multiple_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array([0.1, 0.2]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 5, 2) + + +@pytest.mark.slow +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnShotVectorHessian: + """Class to test the shape of the Hessian with different return types and shot vectors.""" + + def test_hessian_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies) + + def test_hessian_expval_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack(res) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 2, 2) + + def test_hessian_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies) + + def test_hessian_var_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack(res) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 2, 2) + + def test_hessian_probs_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + if diff_method == "spsa": + pytest.skip("SPSA does not support iterated differentiation in Autograd.") + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies, 3) + + def test_hessian_expval_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + if diff_method == "spsa": + pytest.skip("SPSA does not support iterated differentiation in Autograd.") + + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 3, 2, 2) + + +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] + + +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnShotVectorIntegration: + """Tests for the integration of shots with the autograd interface.""" + + def test_single_expectation_value( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """Tests correct output shape and evaluation for a tape + with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = np.array(0.543) + y = np.array(-0.654) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + all_res = qml.jacobian(cost, argnum=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + assert len(all_res) == 2 + + expected = np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]) + tol = TOLS[diff_method] + + for res, exp in zip(all_res, expected): + assert isinstance(res, np.ndarray) + assert res.shape == (num_copies,) + assert np.allclose(res, exp, atol=tol, rtol=0) + + def test_prob_expectation_values( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = np.array(0.543) + y = np.array(-0.654) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + all_res = qml.jacobian(cost, argnum=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + assert len(all_res) == 2 + + expected = np.array( + [ + [ + -np.sin(x), + -(np.cos(y / 2) ** 2 * np.sin(x)) / 2, + -(np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.cos(y / 2) ** 2 * np.sin(x)) / 2, + ], + [ + 0, + -(np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.sin(x / 2) ** 2 * np.sin(y)) / 2, + -(np.sin(x / 2) ** 2 * np.sin(y)) / 2, + ], + ] + ) + + tol = TOLS[diff_method] + + for res, exp in zip(all_res, expected): + assert isinstance(res, np.ndarray) + assert res.shape == (num_copies, 5) + assert np.allclose(res, exp, atol=tol, rtol=0) diff --git a/tests/interfaces/legacy_devices_integration/test_execute_legacy.py b/tests/interfaces/legacy_devices_integration/test_execute_legacy.py new file mode 100644 index 00000000000..e12c5d8ff55 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_execute_legacy.py @@ -0,0 +1,28 @@ +# Copyright 2023 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Interface independent tests for qml.execute +""" + +import pytest + +import pennylane as qml + + +def test_old_interface_no_device_jacobian_products(): + """Test that an error is always raised for the old device interface if device jacobian products are requested.""" + dev = qml.device("default.qubit.legacy", wires=2) + tape = qml.tape.QuantumScript([qml.RX(1.0, wires=0)], [qml.expval(qml.PauliZ(0))]) + with pytest.raises(qml.QuantumFunctionError): + qml.execute((tape,), dev, device_vjp=True) diff --git a/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py new file mode 100644 index 00000000000..c3fd7df5d1a --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py @@ -0,0 +1,901 @@ +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the JAX-JIT interface""" +import numpy as np + +# pylint: disable=protected-access,too-few-public-methods +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + + +class TestJaxExecuteUnitTests: + """Unit tests for jax execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + jax.grad(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_gradients_on_execution(self): + """Test that an error is raised if an gradient transform + is used with grad_on_execution=True""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + grad_on_execution=True, + )[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + jax.grad(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + interface="None", + )[0] + + with pytest.raises(ValueError, match="Unknown interface"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={ + "method": "adjoint_jacobian", + "use_device_state": True, + }, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + jax.jit(cost)(a) + + # adjoint method only performs a single device execution + # gradients are not requested when we only want the results + assert dev.num_executions == 1 + spy.assert_not_called() + + # when the jacobian is requested, we always calculate it at the same time as the results + jax.grad(jax.jit(cost))(a) + spy.assert_called() + + def test_no_gradients_on_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + jax.jit(cost)(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + jax.grad(jax.jit(cost))(a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cachesize=cachesize, + )[0] + + params = jax.numpy.array([0.1, 0.2]) + jax.jit(jax.grad(cost), static_argnums=1)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + custom_cache = {} + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_custom_cache_multiple(self, mocker): + """Test the use of a custom cache object with multiple tapes""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + def cost(a, b, cache): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + res = execute( + [tape1, tape2], + dev, + gradient_fn=param_shift, + cache=cache, + ) + return res[0] + + custom_cache = {} + jax.grad(cost)(a, b, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching produces the optimum number of evaluations.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + def test_caching_adjoint_backward(self): + """Test that caching produces the optimum number of adjoint evaluations + when mode=backward""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one per output dimension + # on the backward pass. + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, also 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestJaxExecuteIntegration: + """Test the jax interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute([tape1, tape2], dev, **execute_kwargs) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = jax.numpy.array(0.1) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = jax.jit(jax.grad(cost))(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + assert tape.trainable_params == [0, 1] + + def cost(a, b): + # An explicit call to _update() is required here to update the + # trainable parameters in between tape executions. + # This is different from how the autograd interface works. + # Unless the update is issued, the validation check related to the + # number of provided parameters fails in the tape: (len(params) != + # required_length) and the tape produces incorrect results. + tape._update() + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return execute([new_tape], dev, **execute_kwargs)[0] + + jac_fn = jax.jit(jax.grad(cost)) + jac = jac_fn(a, b) + + a = jax.numpy.array(0.54) + b = jax.numpy.array(0.8) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = jax.jit(jax.grad(lambda a, b: cost(2 * a, b))) + jac = jac_fn(a, b) + expected = -2 * np.sin(2 * a) + assert np.allclose(jac, expected, atol=tol, rtol=0) + + def test_grad_with_backward_mode(self, execute_kwargs): + """Test jax grad for adjoint diff method in backward mode""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + cost = jax.jit(cost) + + results = jax.grad(cost)(params, cache=None) + for r, e in zip(results, expected_results): + assert jax.numpy.allclose(r, e, atol=1e-7) + + def test_classical_processing_single_tape(self, execute_kwargs): + """Test classical processing within the quantum tape for a single tape""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + c = jax.numpy.array(0.3) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.jit(jax.grad(cost, argnums=(0, 1, 2)), static_argnums=3)(a, b, c, device=dev) + assert len(res) == 3 + + def test_classical_processing_multiple_tapes(self, execute_kwargs): + """Test classical processing within the quantum tape for multiple + tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1] - 7 * result[1] + + res = jax.jit(jax.grad(cost_fn))(params) + assert res.shape == (2,) + + def test_multiple_tapes_output(self, execute_kwargs): + """Test the output types for the execution of multiple quantum tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + + res = jax.jit(cost_fn)(params) + assert isinstance(res, TensorLike) + assert all(isinstance(r, jax.numpy.ndarray) for r in res) + assert all(r.shape == () for r in res) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the jax interface works correctly + with a matrix parameter""" + a = jax.numpy.array(0.1) + U = jax.numpy.array([[0, 1], [1, 0]]) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape.trainable_params = [0] + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.jit(cost, static_argnums=2)(a, U, device=dev) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = jax.grad(cost, argnums=0) + res = jac_fn(a, U, device=dev) + assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def expand(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + qscript = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + qscript = qscript.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return execute([qscript], device, **execute_kwargs)[0] + + a = jax.numpy.array(0.1) + p = jax.numpy.array([0.1, 0.2, 0.3]) + + dev = qml.device("default.qubit.legacy", wires=1) + res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = jax.jit(jax.grad(cost_fn, argnums=1), static_argnums=2) + res = jac_fn(a, p, device=dev) + expected = jax.numpy.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestVectorValuedJIT: + """Test vector-valued returns for the JAX-JIT interface.""" + + @pytest.mark.parametrize( + "ret_type, shape, expected_type", + [ + ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (), tuple), + ([qml.probs(wires=[0, 1])], (4,), jax.numpy.ndarray), + ([qml.probs()], (4,), jax.numpy.ndarray), + ], + ) + def test_shapes(self, execute_kwargs, ret_type, shape, expected_type): + """Test the shape of the result of vector-valued QNodes.""" + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + for r in ret_type: + qml.apply(r) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(cost)(params, cache=None) + assert isinstance(res, expected_type) + + if expected_type is tuple: + for r in res: + assert r.shape == shape + else: + assert res.shape == shape + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + ret_and_output_dim = [ + ([qml.probs(wires=0)], (2,), jax.numpy.ndarray), + ([qml.state()], (4,), jax.numpy.ndarray), + ([qml.density_matrix(wires=0)], (2, 2), jax.numpy.ndarray), + # Multi measurements + ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (), tuple), + ([qml.var(qml.PauliZ(0)), qml.var(qml.PauliZ(1))], (), tuple), + ([qml.probs(wires=0), qml.probs(wires=1)], (2,), tuple), + ] + + @pytest.mark.parametrize("ret, out_dim, expected_type", ret_and_output_dim) + def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim, expected_type): + """Tests the shape of vector-valued QNode results.""" + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + grad_meth = ( + execute_kwargs["gradient_kwargs"]["method"] + if "gradient_kwargs" in execute_kwargs + else "" + ) + if "adjoint" in grad_meth and any( + r.return_type + in (qml.measurements.Probability, qml.measurements.State, qml.measurements.Variance) + for r in ret + ): + pytest.skip("Adjoint does not support probs") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + + for r in ret: + qml.apply(r) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + res = jax.jit(cost, static_argnums=1)(params, cache=None) + + assert isinstance(res, expected_type) + if expected_type is tuple: + for r in res: + assert r.shape == out_dim + else: + assert res.shape == out_dim + + def test_qnode_sample(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + grad_meth = ( + execute_kwargs["gradient_kwargs"]["method"] + if "gradient_kwargs" in execute_kwargs + else "" + ) + if "adjoint" in grad_meth or "backprop" in grad_meth: + pytest.skip("Adjoint does not support probs") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.sample(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q, shots=dev.shots) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + res = jax.jit(cost, static_argnums=1)(params, cache=None) + assert res.shape == (dev.shots,) + + def test_multiple_expvals_grad(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res[0] + res[1] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): + """Test the jacobian computation with multiple tapes with probability + and expectation value computations.""" + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return qml.execute([tape1, tape2], device, **ek, interface=interface)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + res = cost(x, y, dev, interface="jax-jit", ek=execute_kwargs) + + exp = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + for r, e in zip(res, exp): + assert jax.numpy.allclose(r, e, atol=1e-7) + + +def test_diff_method_None_jit(): + """Test that jitted execution works when `gradient_fn=None`.""" + + dev = qml.device("default.qubit.jax", wires=1, shots=10) + + @jax.jit + def wrapper(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return qml.execute([tape], dev, gradient_fn=None) + + assert jax.numpy.allclose(wrapper(jax.numpy.array(0.0))[0], 1.0) diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py similarity index 68% rename from tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py index 065403f9ff3..f109ddc0a4c 100644 --- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,35 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-JIT interface with a QNode""" -import copy - # pylint: disable=too-many-arguments,too-few-public-methods -from functools import partial - import pytest -from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit -# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", True, True], - [ParamShiftDerivativesDevice(), "device", False, True], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "parameter-shift", False, False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_and_qubit_device_and_diff_method = [ ["auto"] + inner_list for inner_list in qubit_device_and_diff_method @@ -57,33 +43,33 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, interface, device_vjp - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") + if diff_method == "backprop": + pytest.skip("Test does not support backprop") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = jax.numpy.array(0.1) + a = np.array(0.1, requires_grad=True) jax.jit(circuit)(a) assert circuit.interface == interface @@ -96,9 +82,7 @@ def circuit(a): assert isinstance(grad, jax.Array) assert grad.shape == () - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -107,12 +91,13 @@ def test_changing_trainability( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + dev = qml.device(dev_name, wires=2) + @qnode( dev, interface=interface, diff_method="parameter-shift", grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -145,21 +130,21 @@ def circuit(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -174,24 +159,21 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the jax interface works correctly with a matrix parameter""" - - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -204,19 +186,16 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that operation and nested tape expansion is differentiable""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = { + "sampler_rng": SEED_FOR_SPSA, + "num_directions": 10, + } tol = TOL_FOR_SPSA class U3(qml.U3): @@ -228,6 +207,13 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -236,7 +222,6 @@ def decomposition(self): diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, p): @@ -263,12 +248,15 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") + a = np.array([0.1, 0.2], requires_grad=True) + dev = qml.device(dev_name, wires=1) + @qnode( dev, interface=interface, @@ -276,7 +264,6 @@ def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, h=1e-8, approx_order=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -291,37 +278,36 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -363,30 +349,29 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") - gradient_kwargs = {} + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -427,21 +412,22 @@ def circuit(a, b): assert r.shape == () assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with a single prob output""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -450,7 +436,6 @@ def test_diff_single_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -480,21 +465,22 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -503,7 +489,6 @@ def test_diff_multi_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -566,21 +551,22 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -589,7 +575,6 @@ def test_diff_expval_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -643,19 +628,23 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") kwargs = {} - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -664,7 +653,6 @@ def test_diff_expval_probs_sub_argnums( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -697,21 +685,19 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "hadamard": + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var") elif diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -720,7 +706,6 @@ def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, i diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -781,8 +766,8 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test device works with diff_method=None.""" - dev = DefaultQubit() + """Test jax device works with diff_method=None.""" + dev = qml.device("default.qubit.jax", wires=1, shots=10) @jax.jit @qml.qnode(dev, diff_method=None, interface=interface) @@ -792,40 +777,52 @@ def circuit(x): assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - @pytest.mark.skip("jax.jit does not work with sample") - def test_changing_shots(self, interface): + def test_changing_shots(self, interface, mocker, tol): """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - assert res.shape == (100, 2) # pylint:disable=comparison-with-callable + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # no additional calls def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" + dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - jit_cost_fn = jax.jit(cost_fn, static_argnames=["shots"]) - res = jax.grad(jit_cost_fn, argnums=[0, 1])(a, b, shots=30000) + # TODO: jit when https://github.com/PennyLaneAI/pennylane/issues/3474 is resolved + res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) + assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) @@ -833,79 +830,44 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface=interface) + # We're choosing interface="jax" such that backprop can be used in the + # test later + @qnode(dev, interface="jax") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" - @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) - def test_shot_vectors_single_measurements(self, interface, shots): - """Test jax-jit can work with shot vectors.""" - - dev = qml.device("default.qubit", shots=shots, seed=4747) - - @jax.jit - @qml.qnode(dev, interface=interface, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.var(qml.PauliZ(0)) - - res = circuit(0.5) - expected = 1 - np.cos(0.5) ** 2 - assert qml.math.allclose(res[0], expected, atol=1e-2) - assert qml.math.allclose(res[1], expected, atol=3e-2) - - g = jax.jacobian(circuit)(0.5) - - expected_g = 2 * np.cos(0.5) * np.sin(0.5) - assert qml.math.allclose(g[0], expected_g, atol=2e-2) - assert qml.math.allclose(g[1], expected_g, atol=2e-2) - - @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) - def test_shot_vectors_multiple_measurements(self, interface, shots): - """Test jax-jit can work with shot vectors.""" - - dev = qml.device("default.qubit", shots=shots, seed=987548) - - @jax.jit - @qml.qnode(dev, interface=interface, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.probs(wires=0) - - res = circuit(0.5) - assert qml.math.allclose(res[0][0], np.cos(0.5), atol=5e-3) - assert qml.math.allclose(res[1][0], np.cos(0.5), atol=5e-3) - expected_probs = np.array([np.cos(0.25) ** 2, np.sin(0.25) ** 2]) - assert qml.math.allclose(res[0][1], expected_probs, atol=5e-3) - assert qml.math.allclose(res[1][1], expected_probs, atol=5e-3) + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): """Test sampling works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -913,19 +875,17 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interfa if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode( + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.sample(qml.Z(0)), qml.sample(qml.s_prod(2, qml.X(0) @ qml.Y(1))) + return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = jax.jit(circuit, static_argnames="shots")(shots=10) + res = jax.jit(circuit)() assert isinstance(res, tuple) @@ -934,7 +894,7 @@ def circuit(): assert isinstance(res[1], jax.Array) assert res[1].shape == (10,) - def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_counts(self, dev_name, diff_method, grad_on_execution, interface): """Test counts works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -942,12 +902,10 @@ def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") + dev = qml.device(dev_name, wires=2, shots=10) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(): qml.Hadamard(wires=[0]) @@ -958,9 +916,9 @@ def circuit(): with pytest.raises( NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): - jax.jit(circuit, static_argnames="shots")(shots=10) + jax.jit(circuit)() else: - res = jax.jit(circuit, static_argnames="shots")(shots=10) + res = jax.jit(circuit)() assert isinstance(res, tuple) @@ -969,32 +927,28 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution, interface): """Test that the gradient of chained QNodes works without error""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) @@ -1020,76 +974,35 @@ def cost(weights): assert len(res) == 2 - def test_postselection_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - elif dev.name == "lightning.qubit": - pytest.xfail("lightning qubit does not support postselection.") - - @qml.qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - phi = jax.numpy.array(1.23) - theta = jax.numpy.array(4.56) - - assert np.allclose(jax.jit(circuit)(phi, theta), jax.jit(expected_circuit)(theta)) - - gradient = jax.jit(jax.grad(circuit, argnums=[0, 1]))(phi, theta) - exp_theta_grad = jax.jit(jax.grad(expected_circuit))(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) - @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test second derivative calculation of a scalar-valued QNode""" + gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: - pytest.skip("Adjoint does not support second derivatives.") + if diff_method == "adjoint": + pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 - gradient_kwargs["h"] = H_FOR_SPSA + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "num_directions": 10} tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1120,25 +1033,31 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { "h": H_FOR_SPSA, - "num_directions": 40, + "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1172,12 +1091,10 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1187,12 +1104,18 @@ def test_hessian_vector_valued( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1236,11 +1159,11 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, interface, device_vjp, grad_on_execution, tol + self, dev_name, diff_method, interface, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1250,12 +1173,18 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1302,11 +1231,11 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1316,13 +1245,19 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -1347,7 +1282,6 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) - hess = jax.jit(jax.jacobian(jac_fn, argnums=[0, 1]))(a, b) expected_hess = np.array( @@ -1367,24 +1301,23 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 - if dev.name == "lightning.qubit" and diff_method == "adjoint": - pytest.xfail("lightning.qubit does not support adjoint with the state.") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) - if not dev.wires: - dev = copy.copy(dev) - dev._wires = qml.wires.Wires([0, 1]) # pylint:disable=protected-access @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1394,7 +1327,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") + assert res.dtype is np.dtype("complex128") # pylint:disable=no-member probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1408,13 +1341,9 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("Adjoint does not support projectors") elif diff_method == "hadamard": @@ -1423,6 +1352,7 @@ def test_projector( gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1431,7 +1361,6 @@ def test_projector( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1454,9 +1383,97 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) +# TODO: Add CV test when return types and custom diff are compatible +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +@pytest.mark.parametrize("interface", ["auto", "jax-jit", "jax"]) +class TestCV: + """Tests for CV integration""" + + def test_first_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = 0.543 + phi = -0.654 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(r, phi) + expected = np.array( + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_second_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = 0.12 + a = 0.765 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(n, a) + expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +# TODO: add support for fwd grad_on_execution to JAX-JIT +@pytest.mark.parametrize("interface", ["auto", "jax-jit"]) +def test_adjoint_reuse_device_state(mocker, interface): + """Tests that the jax interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + jax.grad(circ)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1464,13 +1481,20 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + class PhaseShift(qml.PhaseShift): grad_method = None @@ -1483,7 +1507,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1494,19 +1517,16 @@ def circuit(x, y): x = jax.numpy.array(0.5) y = jax.numpy.array(0.7) circuit(x, y) - jax.grad(circuit, argnums=[0])(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp, interface, mocker, tol + self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradients transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "hadamard": @@ -1519,17 +1539,16 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - spy = mocker.spy(qml.transforms, "hamiltonian_expand") + dev = qml.device(dev_name, wires=3, shots=None) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] - @jax.jit @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1573,16 +1592,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method in ("adjoint", "backprop", "finite-diff"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "hadamard": @@ -1591,6 +1608,8 @@ def test_hamiltonian_finite_shots( gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "h": H_FOR_SPSA, "num_directions": 20} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1599,7 +1618,6 @@ def test_hamiltonian_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1615,12 +1633,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1645,21 +1664,27 @@ def circuit(data, weights, coeffs): # assert np.allclose(grad2_w_c, expected, atol=tol) def test_vmap_compared_param_broadcasting( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not yet support Hamiltonians") + elif diff_method == "hadamard": + pytest.skip("The Hadamard method does not yet support Hamiltonians") + + if diff_method == "backprop": + pytest.skip( + "The backprop method does not yet support parameter-broadcasting with Hamiltonians" + ) + phys_qubits = 2 + if diff_method == "hadamard": + phys_qubits = 3 n_configs = 5 pars_q = np.random.rand(n_configs, 2) + dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1667,80 +1692,42 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) + op = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) + return qml.expval(op) res = _measure_operator() return res - res1 = jax.jit(minimal_circ)(pars_q) - res2 = jax.jit(jax.vmap(minimal_circ))(pars_q) - assert np.allclose(res1, res2, tol) + assert np.allclose( + jax.jit(minimal_circ)(pars_q), jax.jit(jax.vmap(minimal_circ))(pars_q), tol + ) def test_vmap_compared_param_broadcasting_multi_output( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled and when multiple output values are returned.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not yet support Hamiltonians") + elif diff_method == "hadamard": + pytest.skip("The Hadamard method does not yet support Hamiltonians") - n_configs = 5 - pars_q = np.random.rand(n_configs, 2) - - def minimal_circ(params): - @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - cache=None, + if diff_method == "backprop": + pytest.skip( + "The backprop method does not yet support parameter-broadcasting with Hamiltonians" ) - def _measure_operator(): - qml.RY(params[..., 0], wires=0) - qml.RY(params[..., 1], wires=1) - return qml.expval(qml.Z(0) @ qml.Z(1)), qml.expval(qml.X(0) @ qml.X(1)) - - res = _measure_operator() - return res - - res1, res2 = jax.jit(minimal_circ)(pars_q) - vres1, vres2 = jax.jit(jax.vmap(minimal_circ))(pars_q) - assert np.allclose(res1, vres1, tol) - assert np.allclose(res2, vres2, tol) - - def test_vmap_compared_param_broadcasting_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): - """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when - vectorized=True is specified for the callback when caching is disabled and when multiple output values - are returned.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") - elif dev.name == "lightning.qubit" and diff_method == "adjoint": - pytest.xfail("lightning adjoign cannot differentiate probabilities.") - interface = "jax-jit" + phys_qubits = 2 n_configs = 5 pars_q = np.random.rand(n_configs, 2) + dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1748,13 +1735,14 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - return qml.probs(wires=0), qml.probs(wires=1) + op1 = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) + op2 = qml.Hamiltonian([1.0], [qml.PauliX(0) @ qml.PauliX(1)]) + return qml.expval(op1), qml.expval(op2) res = _measure_operator() return res @@ -1770,25 +1758,23 @@ def _measure_operator(): @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestJIT: """Test JAX JIT integration with the QNode and automatic resolution of the correct JAX interface variant.""" - def test_gradient( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface - ): + def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface): """Test derivative calculation of a scalar valued QNode""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if device_vjp and jacobian == jax.jacfwd: - pytest.skip("device vjps not compatible with forward diff.") - elif diff_method == "spsa": + if diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1797,7 +1783,6 @@ def test_gradient( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x): @@ -1821,9 +1806,7 @@ def circuit(x): "ignore:Requested adjoint differentiation to be computed with finite shots." ) @pytest.mark.parametrize("shots", [10, 1000]) - def test_hermitian( - self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface - ): + def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface): """Test that the jax device works with qml.Hermitian and jitting even when shots>0. @@ -1831,6 +1814,13 @@ def test_hermitian( to different reasons, hence the parametrization in the test. """ # pylint: disable=unused-argument + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "backprop": pytest.skip("Backpropagation is unsupported if shots > 0.") @@ -1840,11 +1830,7 @@ def test_hermitian( projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1))) @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circ(projector): return qml.expval(qml.Hermitian(projector, wires=range(2))) @@ -1857,20 +1843,23 @@ def circ(projector): ) @pytest.mark.parametrize("shots", [10, 1000]) def test_probs_obs_none( - self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface + self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface ): """Test that the jax device works with qml.probs, a MeasurementProcess that has obs=None even when shots>0.""" # pylint: disable=unused-argument + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method in ["backprop", "adjoint"]: pytest.skip("Backpropagation is unsupported if shots > 0.") @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(): return qml.probs(wires=0) @@ -1881,30 +1870,23 @@ def circuit(): # reason="Non-trainable parameters are not being correctly unwrapped by the interface" # ) def test_gradient_subset( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test derivative calculation of a scalar valued QNode with respect to a subset of arguments""" - if diff_method == "spsa" and not grad_on_execution and not device_vjp: + if diff_method == "spsa" and not grad_on_execution: pytest.xfail(reason="incorrect jacobian results") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if diff_method == "device" and not grad_on_execution and device_vjp: - pytest.xfail(reason="various runtime-related errors") - - if diff_method == "adjoint" and device_vjp and jacobian is jax.jacfwd: - pytest.xfail(reason="TypeError applying forward-mode autodiff.") + if diff_method == "hadamard" and not grad_on_execution: + pytest.xfail(reason="XLA raised wire error") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + dev = qml.device(dev_name, wires=1) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a, wires=0) @@ -1921,17 +1903,20 @@ def circuit(a, b, c): assert np.allclose(g, expected_g, atol=tol, rtol=0) def test_gradient_scalar_cost_vector_valued_qnode( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test derivative calculation of a scalar valued cost function that uses the output of a vector-valued QNode""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - elif jacobian == jax.jacfwd and device_vjp: - pytest.skip("device vjps are not compatible with forward differentiation.") + if diff_method == "adjoint": + pytest.xfail(reason="Adjoint does not support probs.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1941,7 +1926,6 @@ def test_gradient_scalar_cost_vector_valued_qnode( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1974,26 +1958,21 @@ def cost(x, y, idx): assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0) assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0) - # pylint: disable=unused-argument def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test that the JAX-JIT interface works correctly with a matrix parameter""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("device vjps are not compatible with forward differentiation.") - # pylint: disable=unused-argument + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circ(p, U): qml.QubitUnitary(U, wires=0) @@ -2005,7 +1984,7 @@ def circ(p, U): res = jax.jit(circ)(p, U) assert np.allclose(res, -np.cos(p), atol=tol, rtol=0) - jac_fn = jax.jit(jacobian(circ, argnums=0)) + jac_fn = jax.jit(jax.grad(circ, argnums=(0))) res = jac_fn(p, U) assert np.allclose(res, np.sin(p), atol=tol, rtol=0) @@ -2013,31 +1992,27 @@ def circ(p, U): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturn: """Class to test the shape of the Grad/Jacobian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and one param, the gradient is a float.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2046,30 +2021,27 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + grad = jax.jit(jacobian(circuit))(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2079,9 +2051,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - a, b, shots=shots - ) + grad = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -2089,24 +2059,21 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2115,31 +2082,31 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + grad = jax.jit(jacobian(circuit))(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2148,31 +2115,30 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2182,7 +2148,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(jac, tuple) @@ -2193,25 +2159,24 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2219,33 +2184,29 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2253,9 +2214,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - par_0, par_1, shots=shots - ) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) assert isinstance(jac, tuple) @@ -2274,24 +2233,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2300,7 +2255,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2312,29 +2267,23 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of var.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or only diagonal measurements") + + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2342,9 +2291,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - par_0, par_1, shots=shots - ) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2364,26 +2311,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of var.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or only diagonal measurements") + + dev = qml.device(dev_name, wires=2, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2392,7 +2333,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2404,24 +2345,23 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if device_vjp and jacobian == jax.jacfwd: - pytest.skip("device vjp not compatible with forward differentiation.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2430,7 +2370,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2442,24 +2382,23 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2469,7 +2408,7 @@ def circuit(a, b): a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2489,24 +2428,23 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2515,7 +2453,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2536,17 +2474,23 @@ def circuit(a): @pytest.mark.parametrize("hessian", hessian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturnHessian: """Class to test the shape of the Hessian with different return types.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") par_0 = jax.numpy.array(0.1) @@ -2558,7 +2502,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2586,13 +2529,20 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2601,7 +2551,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2615,10 +2564,12 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + dev = qml.device(dev_name, wires=2) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2632,7 +2583,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2660,14 +2610,16 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") + dev = qml.device(dev_name, wires=2) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2676,7 +2628,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2690,10 +2641,17 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") @@ -2707,7 +2665,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2734,15 +2691,22 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_expval_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2751,7 +2715,6 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2770,10 +2733,12 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + dev = qml.device(dev_name, wires=2) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2787,7 +2752,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2814,14 +2778,16 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_var_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") + dev = qml.device(dev_name, wires=2) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2830,7 +2796,6 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2853,9 +2818,15 @@ def circuit(x): @pytest.mark.parametrize("diff_method", ["parameter-shift", "hadamard"]) def test_jax_device_hessian_shots(hessian, diff_method): """The hessian of multiple measurements with a multiple param array return a single array.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 - @partial(jax.jit, static_argnames="shots") - @qml.qnode(DefaultQubit(), diff_method=diff_method, max_diff=2) + dev = qml.device("default.qubit.jax", wires=num_wires, shots=10000) + + @jax.jit + @qml.qnode(dev, diff_method=diff_method, max_diff=2) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -2864,7 +2835,7 @@ def circuit(x): x = jax.numpy.array([1.0, 2.0]) a, b = x - hess = jax.jit(hessian(circuit), static_argnames="shots")(x, shots=10000) + hess = jax.jit(hessian(circuit))(x) expected_hess = [ [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], @@ -2878,33 +2849,28 @@ def circuit(x): @pytest.mark.parametrize("argnums", [0, 1, [0, 1]]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestSubsetArgnums: def test_single_measurement( self, interface, - dev, + dev_name, diff_method, grad_on_execution, - device_vjp, jacobian, argnums, jit_inside, tol, ): """Test single measurement with different diff methods with argnums.""" + + dev = qml.device(dev_name, wires=3) + kwargs = {} - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transform have a different vjp shape convention.") if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) @qml.qnode( dev, @@ -2912,7 +2878,6 @@ def test_single_measurement( diff_method=diff_method, grad_on_execution=grad_on_execution, cache=False, - device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2942,34 +2907,27 @@ def circuit(a, b): def test_multi_measurements( self, interface, - dev, + dev_name, diff_method, grad_on_execution, - device_vjp, jacobian, argnums, jit_inside, tol, ): """Test multiple measurements with different diff methods with argnums.""" - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transform have a different vjp shape convention.") + dev = qml.device(dev_name, wires=3) kwargs = {} if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) @qml.qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2995,72 +2953,3 @@ def circuit(a, b): else: assert np.allclose(jac[0], expected[0], atol=tol) assert np.allclose(jac[1], expected[1], atol=tol) - - -class TestSinglePrecision: - """Tests for compatibility with single precision mode.""" - - # pylint: disable=import-outside-toplevel - def test_type_conversion_fallback(self): - """Test that if the type isn't int, float, or complex, we still have a fallback.""" - from pennylane.workflow.interfaces.jax_jit import _jax_dtype - - assert _jax_dtype(bool) == jax.numpy.dtype(bool) - - @pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) - def test_float32_return(self, diff_method): - """Test that jax jit works when float64 mode is disabled.""" - jax.config.update("jax_enable_x64", False) - - try: - - @jax.jit - @qml.qnode(qml.device("default.qubit"), diff_method=diff_method) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - grad = jax.grad(circuit)(jax.numpy.array(0.1)) - assert qml.math.allclose(grad, -np.sin(0.1)) - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) - - @pytest.mark.parametrize("diff_method", ("adjoint", "finite-diff")) - def test_complex64_return(self, diff_method): - """Test that jax jit works with differentiating the state.""" - jax.config.update("jax_enable_x64", False) - - try: - tol = 2e-2 if diff_method == "finite-diff" else 1e-6 - - @jax.jit - @qml.qnode(qml.device("default.qubit", wires=1), diff_method=diff_method) - def circuit(x): - qml.RX(x, wires=0) - return qml.state() - - j = jax.jacobian(circuit, holomorphic=True)(jax.numpy.array(0.1 + 0j)) - assert qml.math.allclose(j, [-np.sin(0.05) / 2, -np.cos(0.05) / 2 * 1j], atol=tol) - - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) - - def test_int32_return(self): - """Test that jax jit forward execution works with samples and int32""" - - jax.config.update("jax_enable_x64", False) - - try: - - @jax.jit - @qml.qnode(qml.device("default.qubit", shots=10), diff_method=qml.gradients.param_shift) - def circuit(x): - qml.RX(x, wires=0) - return qml.sample(wires=0) - - _ = circuit(jax.numpy.array(0.1)) - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) diff --git a/tests/interfaces/legacy_devices_integration/test_jax_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_legacy.py new file mode 100644 index 00000000000..e9b7fdd8904 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_legacy.py @@ -0,0 +1,876 @@ +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the JAX-Python interface""" +import numpy as np + +# pylint: disable=protected-access,too-few-public-methods +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + + +class TestJaxExecuteUnitTests: + """Unit tests for jax execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + jax.grad(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if an gradient transform + is used with grad_on_execution=True""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + grad_on_execution=True, + )[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + jax.grad(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + interface="None", + )[0] + + with pytest.raises(ValueError, match="Unknown interface"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad_on_execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=2) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + ) + + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.expval(qml.PauliZ(0))], + ) + + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + ) + return execute( + [tape1, tape2, tape3], + dev, + gradient_fn="device", + gradient_kwargs={ + "method": "adjoint_jacobian", + "use_device_state": True, + }, + ) + + a = jax.numpy.array([0.1, 0.2]) + res = cost(a) + + x, y = a + assert np.allclose(res[0][0], np.cos(x) * np.cos(y)) + assert np.allclose(res[0][1], 1) + assert np.allclose(res[1], np.cos(0.5)) + assert np.allclose(res[2], np.cos(x) * np.cos(y)) + + # adjoint method only performs a single device execution per tape, but gets both result and gradient + assert dev.num_executions == 3 + spy.assert_not_called() + + g = jax.jacobian(cost)(a) + spy.assert_called() + expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y)) + assert qml.math.allclose(g[0][0], expected_g) + assert qml.math.allclose(g[0][1], np.zeros(2)) + assert qml.math.allclose(g[1], np.zeros(2)) + assert qml.math.allclose(g[2], expected_g) + + def test_no_grad_on_execution(self, mocker): + """Test that `grad_on_execution=False` uses the `device.execute_and_gradients`.""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + cost(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + jax.grad(cost)(a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cachesize=cachesize, + )[0] + + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + custom_cache = {} + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_custom_cache_multiple(self, mocker): + """Test the use of a custom cache object with multiple tapes""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + def cost(a, b, cache): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + res = execute( + [tape1, tape2], + dev, + gradient_fn=param_shift, + cache=cache, + ) + return res[0] + + custom_cache = {} + jax.grad(cost)(a, b, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching produces the optimum number of evaluations.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + def test_caching_adjoint_backward(self): + """Test that caching produces the optimum number of adjoint evaluations + when no grad on execution.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one per output dimension + # on the backward pass. + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, also 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestJaxExecuteIntegration: + """Test the jax interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute([tape1, tape2], dev, **execute_kwargs) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = jax.numpy.array(0.1) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = jax.grad(cost)(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + + def cost(a, b): + # An explicit call to _update() is required here to update the + # trainable parameters in between tape executions. + # This is different from how the autograd interface works. + # Unless the update is issued, the validation check related to the + # number of provided parameters fails in the tape: (len(params) != + # required_length) and the tape produces incorrect results. + tape._update() + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return execute([new_tape], dev, **execute_kwargs)[0] + + jac_fn = jax.grad(cost) + jac = jac_fn(a, b) + + a = jax.numpy.array(0.54) + b = jax.numpy.array(0.8) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = jax.grad(lambda a, b: cost(2 * a, b)) + jac = jac_fn(a, b) + expected = -2 * np.sin(2 * a) + assert np.allclose(jac, expected, atol=tol, rtol=0) + + def test_grad_with_different_grad_on_execution(self, execute_kwargs): + """Test jax grad for adjoint diff method with different execution kwargs.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + results = jax.grad(cost)(params, cache=None) + for r, e in zip(results, expected_results): + assert jax.numpy.allclose(r, e, atol=1e-7) + + def test_classical_processing_single_tape(self, execute_kwargs): + """Test classical processing within the quantum tape for a single tape""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + c = jax.numpy.array(0.3) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev) + assert len(res) == 3 + + def test_classical_processing_multiple_tapes(self, execute_kwargs): + """Test classical processing within the quantum tape for multiple + tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1] - 7 * result[1] + + res = jax.grad(cost_fn)(params) + assert res.shape == (2,) + + def test_multiple_tapes_output(self, execute_kwargs): + """Test the output types for the execution of multiple quantum tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + + res = cost_fn(params) + assert isinstance(res, TensorLike) + assert all(isinstance(r, jax.numpy.ndarray) for r in res) + assert all(r.shape == () for r in res) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the jax interface works correctly + with a matrix parameter""" + a = jax.numpy.array(0.1) + U = jax.numpy.array([[0, 1], [1, 0]]) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, U, device=dev) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = jax.grad(cost, argnums=0) + res = jac_fn(a, U, device=dev) + assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + with qml.queuing.AnnotatedQueue() as q_tape: + qml.RX(a, wires=0) + U3(*p, wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q_tape) + tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return execute([tape], device, **execute_kwargs)[0] + + a = jax.numpy.array(0.1) + p = jax.numpy.array([0.1, 0.2, 0.3]) + + dev = qml.device("default.qubit.legacy", wires=1) + res = cost_fn(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = jax.grad(cost_fn, argnums=1) + res = jac_fn(a, p, device=dev) + expected = jax.numpy.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.grad(cost)(params, cache=None) + assert res.shape == (3,) + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestVectorValued: + """Test vector-valued jacobian returns for the JAX Python interface.""" + + def test_multiple_expvals(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jacobian(cost)(params, cache=None) + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert res[0].shape == (3,) + assert isinstance(res[0], jax.numpy.ndarray) + + assert res[1].shape == (3,) + assert isinstance(res[1], jax.numpy.ndarray) + + def test_multiple_expvals_single_par(self, execute_kwargs): + """Tests computing multiple expectation values in a tape with a single + trainable parameter.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jacobian(cost)(params, cache=None) + + assert isinstance(res, tuple) + + assert isinstance(res[0], jax.numpy.ndarray) + assert res[0].shape == (1,) + + assert isinstance(res[1], jax.numpy.ndarray) + assert res[1].shape == (1,) + + def test_multi_tape_fwd(self, execute_kwargs): + """Test the forward evaluation of a cost function that uses the output + of multiple tapes that be vector-valued.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.expval(qml.PauliY(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[1], wires=[0]) + qml.RX(x[1], wires=[0]) + qml.RX(-x[1], wires=[0]) + qml.expval(qml.PauliY(0)) + qml.expval(qml.PauliY(1)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1][0] + + expected = -jax.numpy.sin(params[0]) + -jax.numpy.sin(params[1]) + res = cost_fn(params) + assert jax.numpy.allclose(expected, res) + + def test_multi_tape_jacobian(self, execute_kwargs): + """Test the jacobian computation with multiple tapes.""" + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], device, **ek, interface=interface) + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) + exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + assert np.allclose(exec_jax, exec_autograd) + + res = jax.jacobian(cost, argnums=(0, 1))( + x, y, dev, interface="jax-python", ek=execute_kwargs + ) + + import autograd.numpy as anp + + def cost_stack(x, y, device, interface, ek): + return anp.hstack(cost(x, y, device, interface, ek)) + + exp = qml.jacobian(cost_stack, argnum=(0, 1))( + x_, y_, dev, interface="autograd", ek=execute_kwargs + ) + res_0 = jax.numpy.array([res[0][0][0], res[0][1][0], res[1][0][0], res[1][1][0]]) + res_1 = jax.numpy.array([res[0][0][1], res[0][1][1], res[1][0][1], res[1][1][1]]) + + assert np.allclose(res_0, exp[0]) + assert np.allclose(res_1, exp[1]) + + def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): + """Test the jacobian computation with multiple tapes with probability + and expectation value computations.""" + + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], device, **ek, interface=interface) + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) + exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + assert all( + np.allclose(exec_jax[i][j], exec_autograd[i][j]) for i in range(2) for j in range(2) + ) + + res = jax.jacobian(cost, argnums=(0, 1))( + x, y, dev, interface="jax-python", ek=execute_kwargs + ) + + assert isinstance(res, TensorLike) + assert len(res) == 2 + + for r, exp_shape in zip(res, [(), (2,)]): + assert isinstance(r, tuple) + assert len(r) == 2 + assert len(r[0]) == 2 + assert isinstance(r[0][0], jax.numpy.ndarray) + assert r[0][0].shape == exp_shape + assert isinstance(r[0][1], jax.numpy.ndarray) + assert r[0][1].shape == exp_shape + assert len(r[1]) == 2 + assert isinstance(r[1][0], jax.numpy.ndarray) + assert r[1][0].shape == exp_shape + assert isinstance(r[1][1], jax.numpy.ndarray) + assert r[1][1].shape == exp_shape diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py similarity index 73% rename from tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py index 46f0dc77e22..fa88e315155 100644 --- a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,39 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-Python interface with a QNode""" -# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-implicit-booleaness-not-comparison - -from itertools import product - -import numpy as np +# pylint: disable=too-many-arguments,too-few-public-methods,too-many-public-methods import pytest import pennylane as qml +from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit - -device_seed = 42 - -# device, diff_method, grad_on_execution, device_vjp -device_and_diff_method = [ - [DefaultQubit(seed=device_seed), "backprop", True, False], - [DefaultQubit(seed=device_seed), "finite-diff", False, False], - [DefaultQubit(seed=device_seed), "parameter-shift", False, False], - [DefaultQubit(seed=device_seed), "adjoint", True, False], - [DefaultQubit(seed=device_seed), "adjoint", False, False], - [DefaultQubit(seed=device_seed), "adjoint", True, True], - [DefaultQubit(seed=device_seed), "adjoint", False, True], - [DefaultQubit(seed=device_seed), "spsa", False, False], - [DefaultQubit(seed=device_seed), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], +from pennylane.tape import QuantumScript + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] -interface_and_device_and_diff_method = [ - ["auto"] + inner_list for inner_list in device_and_diff_method -] + [["jax"] + inner_list for inner_list in device_and_diff_method] +interface_and_qubit_device_and_diff_method = [ + ["auto"] + inner_list for inner_list in qubit_device_and_diff_method +] + [["jax"] + inner_list for inner_list in qubit_device_and_diff_method] pytestmark = pytest.mark.jax @@ -58,49 +46,46 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, interface, device_vjp - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = jax.numpy.array(0.1) + a = np.array(0.1, requires_grad=True) circuit(a) assert circuit.interface == interface - # jax doesn't set trainable parameters on regular execution - assert circuit.qtape.trainable_params == [] + # the tape is able to deduce trainable parameters + assert circuit.qtape.trainable_params == [0] # gradients should work grad = jax.grad(circuit)(a) assert isinstance(grad, jax.Array) - # the tape is able to deduce trainable parameters - assert circuit.qtape.trainable_params == [0] assert grad.shape == () - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): # pylint:disable=unused-argument + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -109,7 +94,14 @@ def test_changing_trainability( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - @qnode(dev, interface=interface, diff_method="parameter-shift") + dev = qml.device(dev_name, wires=2) + + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -135,18 +127,27 @@ def circuit(a, b): expected = [-np.sin(a) + np.sin(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + # trainability also updates on evaluation + a = np.array(0.54, requires_grad=False) + b = np.array(0.8, requires_grad=True) + circuit(a, b) + assert circuit.qtape.trainable_params == [1] + + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -161,20 +162,21 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the jax interface works correctly with a matrix parameter""" U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -187,32 +189,39 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that operation and nested tape expansion is differentiable""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 + spsa_kwargs = dict( + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), + num_directions=10, + ) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - class U3(qml.U3): # pylint:disable=too-few-public-methods + class U3(qml.U3): def decomposition(self): theta, phi, lam = self.data wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] + with qml.queuing.AnnotatedQueue() as q_tape: + qml.Rot(lam, theta, -lam, wires=wires) + qml.PhaseShift(phi + lam, wires=wires) + + tape = QuantumScript.from_queue(q_tape) + return tape.operations + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -241,16 +250,23 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): # pylint:disable=unused-argument + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = jax.numpy.array([0.1, 0.2]) + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device(dev_name, wires=1) - @qnode(dev, interface=interface, diff_method="finite-diff", h=1e-8, approx_order=2) + @qnode( + dev, + interface=interface, + diff_method="finite-diff", + h=1e-8, + approx_order=2, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -260,31 +276,30 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode(dev, **kwargs) def circuit(a, b): @@ -295,7 +310,7 @@ def circuit(a, b): res = circuit(a, b) - assert circuit.qtape.trainable_params == [] + assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -305,7 +320,6 @@ def circuit(a, b): res = jax.jacobian(circuit, argnums=[0, 1])(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -325,20 +339,11 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA @@ -346,6 +351,13 @@ def test_jacobian_no_evaluate( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -361,10 +373,11 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - assert isinstance(res[0][0], jax.numpy.ndarray) - for i, j in product((0, 1), (0, 1)): - assert res[i][j].shape == () - assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) + for _res, _exp in zip(res, expected): + for r, e in zip(_res, _exp): + assert isinstance(r, jax.numpy.ndarray) + assert r.shape == () + assert np.allclose(r, e, atol=tol, rtol=0) # call the Jacobian with new parameters a = jax.numpy.array(0.6) @@ -377,27 +390,30 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for i, j in product((0, 1), (0, 1)): - assert res[i][j].shape == () - assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) + for _res, _exp in zip(res, expected): + for r, e in zip(_res, _exp): + assert isinstance(r, jax.numpy.ndarray) + assert r.shape == () + assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with a single prob output""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -429,24 +445,24 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -468,11 +484,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) # pylint:disable=comparison-with-callable + assert res[0].shape == (2,) assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (4,) # pylint:disable=comparison-with-callable + assert res[1].shape == (4,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -511,23 +527,24 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -544,11 +561,11 @@ def circuit(x, y): assert len(res) == 2 assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () # pylint:disable=comparison-with-callable + assert res[0].shape == () assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) # pylint:disable=comparison-with-callable + assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -582,28 +599,32 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" kwargs = {} - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) - if diff_method == "adjoint": - x = x + 0j - y = y + 0j - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -612,8 +633,6 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - if "lightning" in dev.name: - pytest.xfail("lightning does not support measuring probabilities with adjoint.") jac = jax.jacobian(circuit, argnums=[0])(x, y) expected = [ @@ -638,24 +657,21 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if diff_method == "hadamard": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -674,11 +690,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () # pylint:disable=comparison-with-callable + assert res[0].shape == () assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) # pylint:disable=comparison-with-callable + assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -718,40 +734,53 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test device works with diff_method=None.""" + """Test jax device works with diff_method=None.""" + dev = qml.device("default.qubit.jax", wires=1, shots=10) - @qml.qnode(DefaultQubit(), diff_method=None, interface=interface) + @qml.qnode(dev, diff_method=None, interface=interface) def circuit(x): qml.RX(x, wires=0) return qml.expval(qml.PauliZ(0)) - assert jax.numpy.allclose(circuit(jax.numpy.array(0.0), shots=10), 1) + assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - def test_changing_shots(self, interface): + def test_changing_shots(self, interface, mocker, tol): """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) # pylint: disable=comparison-with-callable + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # no additional calls def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" + dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -759,40 +788,51 @@ def cost_fn(a, b): return qml.expval(qml.PauliY(1)) res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) + assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) - def test_update_diff_method(self, interface, mocker): + def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" + # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface=interface) + # We're choosing interface="jax" such that backprop can be used in the + # test later + @qnode(dev, interface="jax") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift -@pytest.mark.parametrize("dev,diff_method,grad_on_execution, device_vjp", device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): + def test_sampling(self, dev_name, diff_method, grad_on_execution): """Test sampling works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") @@ -800,48 +840,43 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - @qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) # pylint:disable=comparison-with-callable + assert res[0].shape == (10,) assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) # pylint:disable=comparison-with-callable + assert res[1].shape == (10,) - def test_counts(self, dev, diff_method, grad_on_execution, device_vjp): + def test_counts(self, dev_name, diff_method, grad_on_execution): """Test counts works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") if diff_method == "adjoint": - pytest.skip("Adjoint errors with finite shots") + pytest.skip("Adjoint warns with finite shots") - @qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1)) + return ( + qml.counts(qml.PauliZ(0), all_outcomes=True), + qml.counts(qml.PauliX(1), all_outcomes=True), + ) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) @@ -850,32 +885,25 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp): + def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" - # pylint:disable=too-few-public-methods + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode( - dev, - interface="jax", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) - @qnode( - dev, - interface="jax", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -900,72 +928,37 @@ def cost(weights): assert len(res) == 2 - def test_postselection_differentiation(self, dev, diff_method, grad_on_execution, device_vjp): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") - - @qml.qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - phi = jax.numpy.array(1.23) - theta = jax.numpy.array(4.56) - - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - - gradient = jax.grad(circuit, argnums=[0, 1])(phi, theta) - exp_theta_grad = jax.grad(expected_circuit)(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) - @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test second derivative calculation of a scalar-valued QNode""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - "max_diff": 2, - } - + kwargs = dict( + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + max_diff=2, + ) if diff_method == "adjoint": pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + spsa_kwargs = dict( + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=100, h=0.001 + ) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(x): qml.RY(x[0], wires=0) @@ -994,7 +987,7 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": @@ -1007,12 +1000,18 @@ def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interfac } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1046,15 +1045,12 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": - qml.math.random.seed(42) gradient_kwargs = { "h": H_FOR_SPSA, "num_directions": 20, @@ -1062,12 +1058,18 @@ def test_hessian_vector_valued( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1111,7 +1113,7 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, interface, grad_on_execution, device_vjp, tol + self, dev_name, diff_method, interface, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} @@ -1125,12 +1127,18 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1177,7 +1185,7 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} @@ -1191,12 +1199,18 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1222,7 +1236,6 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) - hess = jax.jacobian(jac_fn, argnums=[0, 1])(a, b) expected_hess = np.array( @@ -1242,21 +1255,23 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning does not support state adjoint differentiation.") + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1266,7 +1281,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") # pylint:disable=no-member + assert res.dtype is np.dtype("complex128") probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1280,19 +1295,18 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} if diff_method == "adjoint": - pytest.skip("adjoint supports all expvals or only diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1301,7 +1315,6 @@ def test_projector( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1325,7 +1338,94 @@ def circuit(x, y): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +@pytest.mark.parametrize("interface", ["jax", "jax-python"]) +class TestCV: + """Tests for CV integration""" + + def test_first_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = 0.543 + phi = -0.654 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(r, phi) + expected = np.array( + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_second_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = 0.12 + a = 0.765 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(n, a) + expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +@pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"]) +def test_adjoint_reuse_device_state(mocker, interface): + """Tests that the jax interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + jax.grad(circ)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + +@pytest.mark.parametrize( + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1333,14 +1433,21 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1352,7 +1459,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1368,7 +1474,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, interface, device_vjp, mocker, tol + self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None @@ -1386,7 +1492,8 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - spy = mocker.spy(qml.transforms, "hamiltonian_expand") + dev = qml.device(dev_name, wires=3, shots=None) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1394,7 +1501,6 @@ def test_hamiltonian_expansion_analytic( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1439,11 +1545,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured (and not expanded) - if there are non-commuting groups and the number of shots is finite + """Test that the Hamiltonian is expanded if there + are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 @@ -1459,7 +1565,8 @@ def test_hamiltonian_finite_shots( } tol = TOL_FOR_SPSA - spy = mocker.spy(qml.transforms, "hamiltonian_expand") + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1467,7 +1574,6 @@ def test_hamiltonian_finite_shots( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1484,13 +1590,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_not_called() + spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1520,24 +1626,27 @@ def circuit(data, weights, coeffs): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) -class TestReturn: # pylint:disable=too-many-public-methods +class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1546,24 +1655,27 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.grad(circuit)(a, shots=shots) + grad = jax.grad(circuit)(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, shots, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1573,7 +1685,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.grad(circuit, argnums=[0, 1])(a, b, shots=shots) + grad = jax.grad(circuit, argnums=[0, 1])(a, b) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -1581,18 +1693,21 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, shots, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1601,14 +1716,14 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.grad(circuit)(a, shots=shots) + grad = jax.grad(circuit)(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1618,12 +1733,15 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1632,14 +1750,14 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1648,12 +1766,15 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1663,7 +1784,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(a, b) assert isinstance(jac, tuple) @@ -1675,7 +1796,7 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1684,12 +1805,15 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1697,29 +1821,26 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params( - self, - dev, - diff_method, - grad_on_execution, - jacobian, - shots, - interface, - device_vjp, + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if device_vjp and jacobian is jax.jacfwd: - pytest.skip("forward pass can't be done with registered vjp.") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1729,7 +1850,6 @@ def test_jacobian_expval_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1737,7 +1857,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(jac, tuple) @@ -1757,22 +1877,21 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if device_vjp and jacobian is jax.jacfwd: - pytest.skip("forward pass can't be done with registered vjp.") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1781,7 +1900,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1794,16 +1913,18 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": - pytest.skip("adjoint supports either all measurements or only diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Test does not supports adjoint because of var.") + elif diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1813,7 +1934,6 @@ def test_jacobian_var_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1821,7 +1941,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1842,22 +1962,20 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Test does not supports adjoint because of var.") + elif diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1866,7 +1984,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1879,22 +1997,24 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1903,7 +2023,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1916,32 +2036,33 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) - jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(a, b) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1962,22 +2083,23 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1986,7 +2108,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1998,12 +2120,19 @@ def circuit(a): assert jac[1].shape == (4, 2) def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2016,7 +2145,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2024,7 +2152,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2042,7 +2170,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2050,6 +2178,13 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2058,7 +2193,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2066,13 +2200,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2081,6 +2215,7 @@ def test_hessian_var_multiple_params( pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -2091,7 +2226,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2099,7 +2233,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2117,7 +2251,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2127,6 +2261,8 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2135,7 +2271,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2143,15 +2278,22 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2169,7 +2311,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2177,7 +2318,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2213,7 +2354,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2223,13 +2364,19 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( dev, interface=interface, diff_method=diff_method, - device_vjp=device_vjp, max_diff=2, grad_on_execution=grad_on_execution, ) @@ -2239,7 +2386,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2251,7 +2398,7 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2261,6 +2408,8 @@ def test_hessian_probs_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = qml.numpy.array(0.1) par_1 = qml.numpy.array(0.2) @@ -2270,7 +2419,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2278,7 +2426,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2314,7 +2462,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_var_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2324,6 +2472,8 @@ def test_hessian_var_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2332,7 +2482,6 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2340,7 +2489,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2352,11 +2501,14 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="jax") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="jax") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py new file mode 100644 index 00000000000..0ef5c8291de --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py @@ -0,0 +1,933 @@ +# Copyright 2022 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration tests for using the jax interface with shot vectors and with a QNode""" +# pylint: disable=too-many-arguments,too-many-public-methods +import pytest +from flaky import flaky + +import pennylane as qml +from pennylane import numpy as np +from pennylane import qnode + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + +all_shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)] + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], +] + +interface_and_qubit_device_and_diff_method = [ + ["jax"] + inner_list for inner_list in qubit_device_and_diff_method +] + +TOLS = { + "finite-diff": 0.3, + "parameter-shift": 1e-2, + "spsa": 0.32, +} + +jacobian_fn = [jax.jacobian, jax.jacrev, jax.jacfwd] + + +@pytest.mark.parametrize("shots", all_shots) +@pytest.mark.parametrize( + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method +) +class TestReturnWithShotVectors: + """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jac_single_measurement_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + assert isinstance(jac, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jac_single_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + assert isinstance(jac, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + assert j[0].shape == () + assert j[1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_param_probs( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single array is returned with the correct + dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_probs_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (4,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_probs_multiple_param_single_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (4, 2) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_expval_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=1, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == () + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_expval_expval_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_var_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == () + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_var_var_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_single_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(jac, tuple) + assert len(j) == 2 + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == () + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == (4,) + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4, 2) + + def test_hessian_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(hess, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], jax.numpy.ndarray) + assert h[0][0].shape == () + assert h[0][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], jax.numpy.ndarray) + assert h[1][0].shape == () + assert h[1][1].shape == () + + def test_hessian_expval_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, jax.numpy.ndarray) + assert h.shape == (2, 2) + + def test_hessian_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], jax.numpy.ndarray) + assert h[0][0].shape == () + assert h[0][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], jax.numpy.ndarray) + assert h[1][0].shape == () + assert h[1][1].shape == () + + def test_hessian_var_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, jax.numpy.ndarray) + assert h.shape == (2, 2) + + def test_hessian_probs_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], tuple) + assert len(h[0][0]) == 2 + assert isinstance(h[0][0][0], jax.numpy.ndarray) + assert h[0][0][0].shape == () + assert isinstance(h[0][0][1], jax.numpy.ndarray) + assert h[0][0][1].shape == () + assert isinstance(h[0][1], tuple) + assert len(h[0][1]) == 2 + assert isinstance(h[0][1][0], jax.numpy.ndarray) + assert h[0][1][0].shape == () + assert isinstance(h[0][1][1], jax.numpy.ndarray) + assert h[0][1][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], tuple) + assert len(h[1][0]) == 2 + assert isinstance(h[1][0][0], jax.numpy.ndarray) + assert h[1][0][0].shape == (2,) + assert isinstance(h[1][0][1], jax.numpy.ndarray) + assert h[1][0][1].shape == (2,) + assert isinstance(h[1][1], tuple) + assert len(h[1][1]) == 2 + assert isinstance(h[1][1][0], jax.numpy.ndarray) + assert h[1][1][0].shape == (2,) + assert isinstance(h[1][1][1], jax.numpy.ndarray) + assert h[1][1][1].shape == (2,) + + def test_hessian_expval_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because second order diff.") + + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], jax.numpy.ndarray) + assert h[0].shape == (2, 2) + + assert isinstance(h[1], jax.numpy.ndarray) + assert h[1].shape == (2, 2, 2) + + def test_hessian_probs_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = qml.numpy.array(0.1) + par_1 = qml.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], tuple) + assert len(h[0][0]) == 2 + assert isinstance(h[0][0][0], jax.numpy.ndarray) + assert h[0][0][0].shape == () + assert isinstance(h[0][0][1], jax.numpy.ndarray) + assert h[0][0][1].shape == () + assert isinstance(h[0][1], tuple) + assert len(h[0][1]) == 2 + assert isinstance(h[0][1][0], jax.numpy.ndarray) + assert h[0][1][0].shape == () + assert isinstance(h[0][1][1], jax.numpy.ndarray) + assert h[0][1][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], tuple) + assert len(h[1][0]) == 2 + assert isinstance(h[1][0][0], jax.numpy.ndarray) + assert h[1][0][0].shape == (2,) + assert isinstance(h[1][0][1], jax.numpy.ndarray) + assert h[1][0][1].shape == (2,) + assert isinstance(h[1][1], tuple) + assert len(h[1][1]) == 2 + assert isinstance(h[1][1][0], jax.numpy.ndarray) + assert h[1][1][0].shape == (2,) + assert isinstance(h[1][1][1], jax.numpy.ndarray) + assert h[1][1][1].shape == (2,) + + def test_hessian_var_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], jax.numpy.ndarray) + assert h[0].shape == (2, 2) + + assert isinstance(h[1], jax.numpy.ndarray) + assert h[1].shape == (2, 2, 2) + + +@pytest.mark.parametrize("shots", all_shots) +class TestReturnShotVectorsDevice: + """Test for shot vectors with device method adjoint_jacobian.""" + + def test_jac_adjoint_fwd_error(self, shots): + """Test that an error is raised for adjoint forward.""" + dev = qml.device("default.qubit.legacy", wires=1, shots=shots) + + with pytest.warns( + UserWarning, match="Requested adjoint differentiation to be computed with finite shots." + ): + + @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=True) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + if isinstance(shots, tuple): + with pytest.raises( + qml.QuantumFunctionError, + match="Adjoint does not support shot vectors.", + ): + jax.jacobian(circuit)(a) + + def test_jac_adjoint_bwd_error(self, shots): + """Test that an error is raised for adjoint backward.""" + dev = qml.device("default.qubit.legacy", wires=1, shots=shots) + + with pytest.warns( + UserWarning, match="Requested adjoint differentiation to be computed with finite shots." + ): + + @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=False) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + with pytest.raises( + qml.QuantumFunctionError, + match="Adjoint does not support shot vectors.", + ): + jax.jacobian(circuit)(a) + + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], +] + +shots_large = [(1000000, 900000, 800000), (1000000, (900000, 2))] + + +@flaky(max_runs=5) +@pytest.mark.parametrize("shots", shots_large) +@pytest.mark.parametrize( + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method +) +class TestReturnShotVectorIntegration: + """Tests for the integration of shots with the Jax interface.""" + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_single_expectation_value( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """Tests correct output shape and evaluation for a tape + with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]]) + all_res = jacobian(circuit, argnums=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(all_res) == num_copies + + for res in all_res: + assert isinstance(res[0], jax.numpy.ndarray) + assert res[0].shape == () + + assert isinstance(res[1], jax.numpy.ndarray) + assert res[1].shape == () + tol = TOLS[diff_method] + assert np.allclose(res, expected, atol=tol, rtol=0) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_prob_expectation_values( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + all_res = jacobian(circuit, argnums=[0, 1])(x, y) + + tol = TOLS[diff_method] + + assert isinstance(all_res, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(all_res) == num_copies + + for res in all_res: + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], tuple) + assert len(res[0]) == 2 + assert np.allclose(res[0][0], -np.sin(x), atol=tol, rtol=0) + assert isinstance(res[0][0], jax.numpy.ndarray) + assert np.allclose(res[0][1], 0, atol=tol, rtol=0) + assert isinstance(res[0][1], jax.numpy.ndarray) + + assert isinstance(res[1], tuple) + assert len(res[1]) == 2 + assert np.allclose( + res[1][0], + [ + -(np.cos(y / 2) ** 2 * np.sin(x)) / 2, + -(np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.cos(y / 2) ** 2 * np.sin(x)) / 2, + ], + atol=tol, + rtol=0, + ) + assert isinstance(res[1][0], jax.numpy.ndarray) + assert np.allclose( + res[1][1], + [ + -(np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.sin(x / 2) ** 2 * np.sin(y)) / 2, + -(np.sin(x / 2) ** 2 * np.sin(y)) / 2, + ], + atol=tol, + rtol=0, + ) + assert isinstance(res[1][1], jax.numpy.ndarray) diff --git a/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py b/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py new file mode 100644 index 00000000000..6e9739a631f --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py @@ -0,0 +1,47 @@ +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for workflow.set_shots +""" + + +import pennylane as qml +from pennylane.measurements import Shots +from pennylane.workflow import set_shots + + +def test_set_with_shots_class(): + """Test that shots can be set on the old device interface with a Shots class.""" + + dev = qml.devices.DefaultQubitLegacy(wires=1) + with set_shots(dev, Shots(10)): + assert dev.shots == 10 + + assert dev.shots is None + + shot_tuples = Shots((10, 10)) + with set_shots(dev, shot_tuples): + assert dev.shots == 20 + assert dev.shot_vector == list(shot_tuples.shot_vector) + + assert dev.shots is None + + +def test_shots_not_altered_if_False(): + """Test a value of False can be passed to shots, indicating to not override + shots on the device.""" + + dev = qml.devices.DefaultQubitLegacy(wires=1) + with set_shots(dev, False): + assert dev.shots is None diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py similarity index 78% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py index c831fce26cc..f5b6e37a85b 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg +# pylint: disable=too-many-arguments,too-few-public-methods,redefined-outer-name import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -27,25 +26,36 @@ shots_and_num_copies = [((1, (5, 2), 10), 4)] shots_and_num_copies_hess = [((10, (5, 1)), 2)] + +kwargs = { + "finite-diff": {"h": 10e-2}, + "parameter-shift": {}, + "spsa": {"h": 10e-2, "num_directions": 30}, +} + qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", {"h": 10e-2}], - [DefaultQubit(), "parameter-shift", {}], - [ - DefaultQubit(), - "spsa", - {"h": 10e-2, "num_directions": 20, "sampler_rng": np.random.default_rng(42)}, - ], + ["default.qubit.legacy", "finite-diff"], + ["default.qubit.legacy", "parameter-shift"], + ["default.qubit.legacy", "spsa"], ] TOLS = { "finite-diff": 0.3, "parameter-shift": 1e-2, - "spsa": 0.5, + "spsa": 0.3, } +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -54,13 +64,14 @@ class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)) @@ -68,7 +79,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -77,13 +88,14 @@ def circuit(a, **_): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)) @@ -92,7 +104,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -104,13 +116,14 @@ def circuit(a, b, **_): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) @@ -118,7 +131,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -127,14 +140,15 @@ def circuit(a, **_): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.probs(wires=[0, 1]) @@ -142,7 +156,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -151,14 +165,15 @@ def circuit(a, **_): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) @@ -167,7 +182,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -179,14 +194,15 @@ def circuit(a, b, **_): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) @@ -194,7 +210,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -203,22 +219,24 @@ def circuit(a, **_): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=1, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -230,13 +248,14 @@ def circuit(x, y, **_): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) qml.RY(a[2], wires=0) @@ -245,7 +264,7 @@ def circuit(a, **_): a = tf.Variable([0.7, 0.9, 1.1], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -254,13 +273,14 @@ def circuit(a, **_): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -268,7 +288,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -277,13 +297,14 @@ def circuit(a, **_): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -292,7 +313,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -304,13 +325,14 @@ def circuit(a, b, **_): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -318,7 +340,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -329,7 +351,7 @@ def circuit(a, **_): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -338,7 +360,7 @@ class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" @@ -346,12 +368,14 @@ def test_hessian_expval_multiple_params( # TODO: Find out why. pytest.skip("SPSA gradient does not support this particular test case") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) @@ -359,7 +383,7 @@ def circuit(x, y, **_): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -374,11 +398,11 @@ def circuit(x, y, **_): assert h.shape == (2, num_copies) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((20000, 18000, 16000), 3), ((20000, (18000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -387,23 +411,24 @@ class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -420,23 +445,24 @@ def circuit(x, y, **_): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py new file mode 100644 index 00000000000..057f93d0038 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py @@ -0,0 +1,1028 @@ +# Copyright 2018-2021 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the TensorFlow interface""" +# pylint: disable=protected-access,too-few-public-methods +import numpy as np +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import finite_diff, param_shift + +pytestmark = pytest.mark.tf + +tf = pytest.importorskip("tensorflow", minversion="2.1") + + +class TestTensorFlowExecuteUnitTests: + """Unit tests for TensorFlow execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = tf.Variable([0.1, 0.2], dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + interface="tf", + )[0] + + res = t.jacobian(res, a) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if a gradient transform + is used with grad_on_execution""" + a = tf.Variable([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + execute([tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface="tf") + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable([0.1, 0.2]) + spy = mocker.spy(dev, "execute_and_gradients") + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface="tf", + ) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_no_grad_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + a = tf.Variable([0.1, 0.2]) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface="tf", + )[0] + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + t.jacobian(res, a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + a = tf.Variable([0.1, 0.2]) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, gradient_fn=param_shift, cachesize=2, interface="tf")[0] + + t.jacobian(res, a) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + a = tf.Variable([0.1, 0.2]) + custom_cache = {} + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, gradient_fn=param_shift, cache=custom_cache, interface="tf")[ + 0 + ] + + t.jacobian(res, a) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + unwrapped_tape = qml.transforms.convert_to_numpy_parameters(tape)[0][0] + h = unwrapped_tape.hash + + assert h in cache + assert np.allclose(cache[h], res) + + def test_caching_param_shift(self): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum.""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="tf")[0] + + # Without caching, and non-vectorized, 9 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + with tf.GradientTape(persistent=True) as t: + res = cost(a, cache=None) + t.jacobian(res, a, experimental_use_pfor=False) + assert dev.num_executions == 9 + + # With caching, and non-vectorized, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + with tf.GradientTape(persistent=True) as t: + res = cost(a, cache=True) + t.jacobian(res, a) + assert dev.num_executions == 5 + + # In vectorized mode, 5 evaluations are required to compute + # the Jacobian regardless of caching: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + with tf.GradientTape() as t: + res = cost(a, cache=None) + t.jacobian(res, a) + assert dev.num_executions == 5 + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = tf.Variable(np.arange(1, num_params + 1) / 10, dtype=tf.float64) + + N = params.shape[0] + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], dev, gradient_fn=param_shift, cache=cache, interface="tf", max_diff=2 + )[0] + + # No caching: number of executions is not ideal + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + res = cost(params, cache=False) + grad = t1.gradient(res, params) + hess1 = t2.jacobian(grad, params) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params * 1.0 + expected = np.array( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + nonideal_runs = dev.num_executions + + # Use caching: number of executions is ideal + dev._num_executions = 0 + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + res = cost(params, cache=True) + grad = t1.gradient(res, params) + hess2 = t2.jacobian(grad, params) + + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < nonideal_runs + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift, "interface": "tf"}, + {"gradient_fn": param_shift, "interface": "auto"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "tf", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "tf", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "auto", + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestTensorFlowExecuteIntegration: + """Test the TensorFlow interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable(0.1) + b = tf.Variable(0.2) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + res = execute([tape1, tape2], dev, **execute_kwargs) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + assert isinstance(res[0], tf.Tensor) + assert isinstance(res[1], tf.Tensor) + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = tf.Variable(0.1, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + + res = t.jacobian(res, a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_jacobian(self, execute_kwargs, tol): + """Test jacobian calculation""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, max_diff=2, **execute_kwargs)[0] + res = tf.stack(res) + + expected = [np.cos(a), -np.cos(a) * np.sin(b)] + assert np.allclose(res, expected, atol=tol, rtol=0) + + (agrad, bgrad) = t.jacobian(res, [a, b]) + assert agrad.shape == (2,) + assert bgrad.shape == (2,) + + expected = [[-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]] + assert np.allclose(expected, [agrad, bgrad], atol=tol, rtol=0) + + def test_tape_no_parameters(self, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + dev = qml.device("default.qubit.legacy", wires=1) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = 1.0 * params + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(0.5, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) + res = tf.stack(res) + + expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t.gradient(res, params) + expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] + assert np.allclose(grad, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + res = execute([tape], dev, **execute_kwargs)[0] + res = tf.stack(res) + + t.jacobian(res, [a, b]) + + a = tf.Variable(0.54, dtype=tf.float64) + b = tf.Variable(0.8, dtype=tf.float64) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + res2 = tf.stack(res2) + + expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac2 = t.jacobian(res2, [a, b]) + expected = [ + [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], + [0, -tf.cos(2 * a) * tf.cos(b)], + ] + assert np.allclose(jac2, expected, atol=tol, rtol=0) + + def test_reusing_pre_constructed_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape that was previously constructed + *outside of* a gradient tape, by passing new parameters""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([a, b], [0, 1]) + assert tape.trainable_params == [0, 1] + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + t.jacobian(res, [a, b]) + + a = tf.Variable(0.54, dtype=tf.float64) + b = tf.Variable(0.8, dtype=tf.float64) + + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + res2 = qml.math.stack(res2) + + expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac2 = t.jacobian(res2, [a, b]) + expected = [ + [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], + [0, -tf.cos(2 * a) * tf.cos(b)], + ] + assert np.allclose(jac2, expected, atol=tol, rtol=0) + + def test_classical_processing(self, execute_kwargs): + """Test classical processing within the quantum tape""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.constant(0.2, dtype=tf.float64) + c = tf.Variable(0.3, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + tf.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [0, 2] + assert tape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] + + res = t.jacobian(res, [a, b, c]) + assert isinstance(res[0], tf.Tensor) + assert res[1] is None + assert isinstance(res[2], tf.Tensor) + + def test_no_trainable_parameters(self, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + b = tf.constant(0.2, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(0.2, wires=0) + qml.RX(b, wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + assert res.shape == (2,) + assert isinstance(res, tf.Tensor) + + res = t.jacobian(res, b) + assert res is None + + @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) + def test_matrix_parameter(self, execute_kwargs, U, tol): + """Test that the TF interface works correctly + with a matrix parameter""" + a = tf.Variable(0.1, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1] + + assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) + + res = t.jacobian(res, a) + assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tape expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + dev = qml.device("default.qubit.legacy", wires=1) + a = np.array(0.1) + p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + + with tf.GradientTape() as tape: + with qml.queuing.AnnotatedQueue() as q_qtape: + qml.RX(a, wires=0) + U3(p[0], p[1], p[2], wires=0) + qml.expval(qml.PauliX(0)) + + qtape = qml.tape.QuantumScript.from_queue(q_qtape) + res = execute([qtape], dev, **execute_kwargs)[0] + + expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( + tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, p) + expected = np.array( + [ + tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), + tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) + - tf.sin(p[1]) + * (tf.cos(a) * tf.sin(p[0]) + tf.cos(p[0]) * tf.sin(a) * tf.sin(p[2])), + tf.sin(a) + * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + expected = np.array( + [ + [tf.cos(x / 2) ** 2, tf.sin(x / 2) ** 2], + [(1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = t.jacobian(res, [x, y]) + expected = np.array( + [ + [ + [-tf.sin(x) / 2, tf.sin(x) / 2], + [-tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], + ], + [ + [0, 0], + [-tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = tf.experimental.numpy.hstack(res) + + expected = np.array( + [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = t.jacobian(res, [x, y]) + expected = np.array( + [ + [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], + [0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_sampling(self, execute_kwargs): + """Test sampling works as expected""" + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(tf.Variable(0.1), wires=0) + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + assert res.shape == (2, 10) + assert isinstance(res, tf.Tensor) + + +@pytest.mark.parametrize("interface", ["auto", "tf"]) +class TestHigherOrderDerivatives: + """Test that the TensorFlow execute function can be differentiated""" + + @pytest.mark.slow + @pytest.mark.parametrize( + "params", + [ + tf.Variable([0.543, -0.654], dtype=tf.float64), + tf.Variable([0, -0.654], dtype=tf.float64), + tf.Variable([-2.0, 0], dtype=tf.float64), + ], + ) + def test_parameter_shift_hessian(self, params, tol, interface): + """Tests that the output of the parameter-shift transform + can be differentiated using tensorflow, yielding second derivatives.""" + dev = qml.device("default.qubit.tf", wires=2) + x, y = params * 1.0 + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(params[0], wires=0) + qml.RY(params[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, interface=interface, max_diff=2 + ) + res = result[0] + result[1][0] + + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t1.gradient(res, params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + + hess = t2.jacobian(grad, params) + expected = np.array( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(hess, expected, atol=tol, rtol=0) + + def test_hessian_vector_valued(self, tol, interface): + """Test hessian calculation of a vector valued QNode""" + dev = qml.device("default.qubit.tf", wires=1) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + + with tf.GradientTape() as t2: + with tf.GradientTape(persistent=True) as t1: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], dev, gradient_fn=param_shift, interface=interface, max_diff=2 + )[0] + res = tf.stack(res) + + g = t1.jacobian(res, params, experimental_use_pfor=False) + + hess = t2.jacobian(g, params) + + a, b = params * 1.0 + + expected_res = [ + 0.5 + 0.5 * tf.cos(a) * tf.cos(b), + 0.5 - 0.5 * tf.cos(a) * tf.cos(b), + ] + assert np.allclose(res, expected_res, atol=tol, rtol=0) + + expected_g = [ + [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)], + [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)], + ] + assert np.allclose(g, expected_g, atol=tol, rtol=0) + + expected_hess = [ + [ + [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)], + [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)], + ], + [ + [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)], + [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)], + ], + ] + + np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) + + def test_adjoint_hessian(self, interface): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface=interface, + )[0] + + grad = t1.gradient(res, params) + assert grad is not None + assert grad.dtype == tf.float64 + assert grad.shape == params.shape + + hess = t2.jacobian(grad, params) + assert hess is None + + def test_max_diff(self, tol, interface): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.tf", wires=2) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + x, y = params * 1.0 + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(params[0], wires=0) + qml.RY(params[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface=interface + ) + res = result[0] + result[1][0] + + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t1.gradient(res, params) + + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + + hess = t2.jacobian(grad, params) + assert hess is None + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift, "interface": "tensorflow"}, + {"gradient_fn": finite_diff, "interface": "tensorflow"}, + {"gradient_fn": param_shift, "interface": "auto"}, + {"gradient_fn": finite_diff, "interface": "auto"}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + return tf.stack(execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1.numpy() + d = coeffs2.numpy()[0] + x, y = weights.numpy() + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian_expected(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1.numpy() + d = coeffs2.numpy()[0] + x, y = weights.numpy() + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) + coeffs2 = tf.constant([0.7], dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as tape: + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, [weights, coeffs1, coeffs2]) + expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) + assert res[1] is None + assert res[2] is None + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + coeffs2 = tf.Variable([0.7], dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as tape: + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, [weights, coeffs1, coeffs2]) + expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) + assert np.allclose(res[1], expected[:, 2:5], atol=tol, rtol=0) + assert np.allclose(res[2], expected[:, 5:], atol=tol, rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py similarity index 70% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py index e3c4597ae06..5b03f9da10f 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2020 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,31 +13,26 @@ # limitations under the License. """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, use-implicit-booleaness-not-comparison import pytest import pennylane as qml from pennylane import qnode -from pennylane.devices import DefaultQubit + +# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison + pytestmark = pytest.mark.tf tf = pytest.importorskip("tensorflow") -# device, diff_method, grad_on_execution, device_vjp + qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=4), "adjoint", False, True], - [qml.device("lightning.qubit", wires=4), "adjoint", False, False], - [qml.device("lightning.qubit", wires=4), "adjoint", True, True], - [qml.device("lightning.qubit", wires=4), "adjoint", True, False], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] TOL_FOR_SPSA = 1.0 @@ -50,25 +45,25 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with TensorFlow integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -89,7 +84,7 @@ def circuit(a): # with the interface, the tape returns tensorflow tensors assert isinstance(res, tf.Tensor) - assert res.shape == () + assert res.shape == tuple() # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -97,20 +92,23 @@ def circuit(a): # gradients should work grad = tape.gradient(res, a) assert isinstance(grad, tf.Tensor) - assert grad.shape == () + assert grad.shape == tuple() - def test_interface_swap(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_interface_swap(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the TF interface can be applied to a QNode with a pre-existing interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -137,44 +135,51 @@ def circuit(a): assert np.allclose(res1, res2, atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_drawing(self, dev_name, diff_method, grad_on_execution, interface): """Test circuit drawing when using the TF interface""" x = tf.Variable(0.1, dtype=tf.float64) y = tf.Variable([0.2, 0.3], dtype=tf.float64) z = tf.Variable(0.4, dtype=tf.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) qml.RY(p2[0] * p2[1], wires=1) qml.RX(kwargs["p3"], wires=0) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) + return qml.state() result = qml.draw(circuit)(p1=x, p3=z) - expected = "0: ──RX(0.10)──RX(0.40)─╭●── \n1: ──RY(0.06)───────────╰X── " + expected = "0: ──RX(0.10)──RX(0.40)─╭●── State\n1: ──RY(0.06)───────────╰X── State" assert result == expected - def test_jacobian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_jacobian(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test jacobian calculation""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) @@ -185,7 +190,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -197,17 +202,61 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [a, b]) expected = [[-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_jacobian_dtype(self, dev_name, diff_method, grad_on_execution, interface): + """Test calculating the jacobian with a different datatype""" + if diff_method == "backprop": + pytest.skip("Test does not support backprop") + + a = tf.Variable(0.1, dtype=tf.float32) + b = tf.Variable(0.2, dtype=tf.float32) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, r_dtype=np.float32) + + @qnode( + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + + with tf.GradientTape() as tape: + res = circuit(a, b) + res = tf.stack(res) + + assert circuit.qtape.trainable_params == [0, 1] + + assert isinstance(res, tf.Tensor) + assert res.shape == (2,) + assert res.dtype is tf.float32 + + res = tape.jacobian(res, [a, b]) + assert [r.dtype is tf.float32 for r in res] + + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting finite-difference jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite diff and spsa.") a = tf.Variable([0.1, 0.2]) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, interface=interface, @@ -215,21 +264,18 @@ def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, approx_order=2, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + tape.jacobian(res, a) - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method in ["backprop", "adjoint", "spsa"]: @@ -238,16 +284,21 @@ def test_changing_trainability( a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) + num_wires = 2 + diff_kwargs = {} - if diff_method == "finite-diff": + if diff_method == "hadamard": + num_wires = 3 + elif diff_method == "finite-diff": diff_kwargs = {"approx_order": 2, "strategy": "center"} + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **diff_kwargs, ) def circuit(a, b): @@ -256,7 +307,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -266,7 +317,7 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, [a, b]) expected = [ [-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)], @@ -277,7 +328,7 @@ def circuit(a, b): a = tf.Variable(0.54, dtype=tf.float64) b = tf.constant(0.8, dtype=tf.float64) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -287,22 +338,25 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) expected = [-tf.sin(a), tf.sin(a) * tf.sin(b)] assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(x, y, z): qml.RY(x * z, wires=0) @@ -310,30 +364,30 @@ def circuit(x, y, z): qml.RX(z + z**2 + tf.sin(a), wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b, c) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [0, 2] assert circuit.qtape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] - res = tape.jacobian(res, [a, b, c], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [a, b, c]) assert isinstance(res[0], tf.Tensor) assert res[1] is None assert isinstance(res[2], tf.Tensor) - def test_no_trainable_parameters( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_no_trainable_parameters(self, dev_name, diff_method, grad_on_execution, interface): """Test evaluation if there are no trainable parameters""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b): qml.RY(a, wires=0) @@ -344,7 +398,7 @@ def circuit(a, b): a = 0.1 b = tf.constant(0.2, dtype=tf.float64) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -355,30 +409,31 @@ def circuit(a, b): assert isinstance(res, tf.Tensor) # can't take the gradient with respect to "a" since it's a Python scalar - grad = tape.jacobian(res, b, experimental_use_pfor=not device_vjp) + grad = tape.jacobian(res, b) assert grad is None @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, U, tol, interface - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, U, tol, interface): """Test that the TF interface works correctly with a matrix parameter""" a = tf.Variable(0.1, dtype=tf.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(U, a): qml.QubitUnitary(U, wires=0) qml.RY(a, wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(U, a) if diff_method == "finite-diff": @@ -386,23 +441,18 @@ def circuit(U, a): assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) - res = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, a) assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that operation and nested tapes expansion is differentiable""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA class U3(qml.U3): @@ -414,6 +464,13 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + a = np.array(0.1) p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) @@ -423,7 +480,7 @@ def circuit(a, p): U3(p[0], p[1], p[2], wires=0) return qml.expval(qml.PauliX(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( @@ -431,7 +488,7 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, p, experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, p) expected = np.array( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), @@ -450,9 +507,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, interface): + def test_changing_shots(self, mocker, tol, interface): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -461,21 +518,31 @@ def circuit(weights): qml.RY(weights[0], wires=0) qml.RX(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(weights) + res = circuit(weights) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg - assert res.shape == (100, 2) + circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(weights) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" # pylint: disable=unexpected-keyword-arg - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -490,6 +557,7 @@ def circuit(weights): res = circuit(weights, shots=[10000, 10000, 10000]) res = tf.transpose(tf.stack(res)) + assert dev.shots is None assert len(res) == 3 jacobian = tape.jacobian(res, weights) @@ -500,7 +568,7 @@ def test_multiple_gradient_integration(self, tol, interface): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -514,7 +582,7 @@ def circuit(weights): with tf.GradientTape() as tape: res1 = circuit(weights) - assert qml.math.shape(res1) == () + assert qml.math.shape(res1) == tuple() res2 = circuit(weights, shots=[(1, 1000)]) # pylint: disable=unexpected-keyword-arg assert qml.math.shape(res2) == (1000,) @@ -525,7 +593,7 @@ def circuit(weights): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=100) weights = tf.Variable([0.543, -0.654], dtype=tf.float64) spy = mocker.spy(qml, "execute") @@ -537,45 +605,116 @@ def circuit(weights): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - circuit(weights, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert circuit.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + circuit(weights) + assert circuit.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used circuit(weights) - assert circuit.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert circuit.gradient_fn is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + circuit(weights, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert circuit.gradient_fn == "backprop" + + circuit(weights) + assert circuit.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + +@pytest.mark.parametrize("interface", ["auto", "tf"]) +class TestAdjoint: + """Specific integration tests for the adjoint method""" + + def test_reuse_state(self, mocker, interface): + """Tests that the TF interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, diff_method="adjoint", interface=interface) + def circ(x): + qml.RX(x[0], wires=0) + qml.RY(x[1], wires=1) + qml.CNOT(wires=(0, 1)) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + weights = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = 1.0 * weights + + with tf.GradientTape() as tape: + res = tf.reduce_sum(circ(weights)) + + grad = tape.gradient(res, weights) + expected_grad = [-tf.sin(x), tf.cos(y)] + + assert np.allclose(grad, expected_grad) + assert circ.device.num_executions == 1 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + + def test_reuse_state_multiple_evals(self, mocker, tol, interface): + """Tests that the TF interface reuses the device state for adjoint differentiation, + even where there are intermediate evaluations.""" + dev = qml.device("default.qubit.legacy", wires=2) + + x_val = 0.543 + y_val = -0.654 + x = tf.Variable(x_val, dtype=tf.float64) + y = tf.Variable(y_val, dtype=tf.float64) + + @qnode(dev, diff_method="adjoint", interface=interface) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + with tf.GradientTape() as tape: + res1 = circuit(x, y) + + assert np.allclose(res1, np.cos(x_val), atol=tol, rtol=0) + + # intermediate evaluation with different values + circuit(tf.math.tan(x), tf.math.cosh(y)) + + # the adjoint method will continue to compute the correct derivative + grad = tape.gradient(res1, x) + assert np.allclose(grad, -np.sin(x_val), atol=tol, rtol=0) + assert dev.num_executions == 2 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + self, dev_name, diff_method, grad_on_execution, tol, interface ): """Tests correct output shape and evaluation for a tape with multiple probs outputs""" - - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") - - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -586,7 +725,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.probs(wires=[0]), qml.probs(wires=[1]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(x, y) res = tf.stack(res) @@ -598,7 +737,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [x, y]) expected = np.array( [ [ @@ -613,26 +752,25 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_ragged_differentiation(self, dev_name, diff_method, grad_on_execution, tol, interface): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") - - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -643,7 +781,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(x, y) res = tf.experimental.numpy.hstack(res) @@ -656,7 +794,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [x, y]) expected = np.array( [ [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], @@ -665,20 +803,24 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -708,18 +850,24 @@ def circuit(x): ] assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_hessian(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -748,20 +896,24 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -804,19 +956,25 @@ def circuit(x): np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + self, dev_name, diff_method, grad_on_execution, tol, interface ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=0) @@ -855,18 +1013,24 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_hessian_ragged(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -917,21 +1081,23 @@ def circuit(x): ] np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_state(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -941,6 +1107,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) + assert res.dtype is tf.complex128 probs = tf.math.abs(res) ** 2 return probs[0] + probs[2] @@ -955,27 +1122,21 @@ 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, dtype - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the variance of a projector is correctly returned""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "hadamard": pytest.skip("Variance not implemented yet.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - P = tf.constant(state, dtype=dtype) + dev = qml.device(dev_name, wires=2) + P = tf.constant(state) x, y = 0.765, -0.654 weights = tf.Variable([x, y], dtype=tf.float64) @@ -1000,69 +1161,95 @@ def circuit(weights): ] assert np.allclose(grad, expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = tf.Variable(0.543, dtype=tf.float64) + phi = tf.Variable(-0.654, dtype=tf.float64) + + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + with tf.GradientTape() as tape: + res = circuit(r, phi) + + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + grad = tape.gradient(res, [r, phi]) + expected = [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + assert np.allclose(grad, expected, atol=tol, rtol=0) - phi = tf.Variable(1.23) - theta = tf.Variable(4.56) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = tf.Variable(0.12, dtype=tf.float64) + a = tf.Variable(0.765, dtype=tf.float64) - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) - with tf.GradientTape() as res_tape: - res = circuit(phi, theta) - gradient = res_tape.gradient(res, [phi, theta]) + with tf.GradientTape() as tape: + res = circuit(n, a) - with tf.GradientTape() as expected_tape: - expected = expected_circuit(theta) - exp_theta_grad = expected_tape.gradient(expected, theta) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) - assert np.allclose(gradient, [0.0, exp_theta_grad]) + # circuit jacobians + grad = tape.gradient(res, [n, a]) + expected = [2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)] + assert np.allclose(grad, expected, atol=tol, rtol=0) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the TF interface""" - def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution, interface): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) class PhaseShift(qml.PhaseShift): grad_method = None @@ -1076,7 +1263,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.Hadamard(wires=0) @@ -1088,6 +1274,7 @@ def circuit(x): with tf.GradientTape() as t2: with tf.GradientTape() as t1: loss = circuit(x) + res = t1.gradient(loss, x) assert np.allclose(res, -3 * np.sin(3 * x)) @@ -1099,13 +1286,20 @@ def circuit(x): @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + class PhaseShift(qml.PhaseShift): grad_method = None @@ -1118,7 +1312,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1136,24 +1329,24 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, tol, interface + self, dev_name, diff_method, grad_on_execution, max_diff, tol, interface ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "max_diff": max_diff, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, + grad_on_execution=grad_on_execution, + max_diff=max_diff, + interface=interface, + ) if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint/hadamard method does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1200,14 +1393,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker, interface ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1220,6 +1413,8 @@ def test_hamiltonian_finite_shots( elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1228,7 +1423,6 @@ def test_hamiltonian_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1244,19 +1438,17 @@ def circuit(data, weights, coeffs): c = tf.Variable([-0.6543, 0.24, 0.54], dtype=tf.float64) # test output - with tf.GradientTape(persistent=True) as _t2: + with tf.GradientTape(persistent=True) as t2: with tf.GradientTape() as t1: - res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients grad = t1.gradient(res, [d, w, c]) - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1266,36 +1458,33 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[2], expected_c, atol=tol) # test second-order derivatives - # TODO: figure out why grad2_c is np.zeros((3,3)) instead of None - # if diff_method == "parameter-shift" and max_diff == 2: - # grad2_c = _t2.jacobian(grad[2], c) - # print(grad2_c, grad[2], c) - # assert grad2_c is None + if diff_method == "parameter-shift" and max_diff == 2: + grad2_c = t2.jacobian(grad[2], c) + assert grad2_c is None - # grad2_w_c = _t2.jacobian(grad[1], c) - # expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ - # 0, - # -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), - # -np.sin(d[1] + w[1]), - # ] - # assert np.allclose(grad2_w_c, expected, atol=tol) + grad2_w_c = t2.jacobian(grad[1], c) + expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ + 0, + -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), + -np.sin(d[1] + w[1]), + ] + assert np.allclose(grad2_w_c, expected, atol=tol) class TestSample: """Tests for the sample integration""" - # pylint:disable=unexpected-keyword-arg - def test_sample_dimension(self): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 @@ -1307,14 +1496,15 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) @@ -1325,67 +1515,76 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" + n_sample = 10 - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=10) + result = circuit() assert len(result) == 3 - assert result[0].shape == (10,) + assert result[0].shape == (n_sample,) assert result[1].shape == () assert result[2].shape == () assert isinstance(result[0], tf.Tensor) assert isinstance(result[1], tf.Tensor) assert isinstance(result[2], tf.Tensor) - assert result[0].dtype is tf.float64 # pylint:disable=no-member - assert result[1].dtype is tf.float64 # pylint:disable=no-member - assert result[2].dtype is tf.float64 # pylint:disable=no-member + assert result[0].dtype is tf.float64 + assert result[1].dtype is tf.float64 + assert result[2].dtype is tf.float64 def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (10,)) + assert np.array_equal(result.shape, (n_sample,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=10) + result = circuit() result = tf.stack(result) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (3, 10)) + assert np.array_equal(result.shape, (3, n_sample)) assert result.dtype == tf.float64 def test_counts(self): """Test counts works as expected for TF""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) - # pylint:disable=unsubscriptable-object,no-member - @qnode(DefaultQubit(), interface="tf") + @qnode(dev, interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)) - res = circuit(shots=100) + res = circuit() assert isinstance(res, dict) assert list(res.keys()) == [-1, 1] @@ -1415,11 +1614,12 @@ class TestAutograph: def test_autograph_gradients(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + @qnode(dev, diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1428,7 +1628,7 @@ def circuit(x, y): with tf.GradientTape() as tape: p0, p1 = circuit(x, y) - loss = p0[0] + p1[1] # pylint:disable=unsubscriptable-object + loss = p0[0] + p1[1] expected = tf.cos(x / 2) ** 2 + (1 - tf.cos(x) * tf.cos(y)) / 2 assert np.allclose(loss, expected, atol=tol, rtol=0) @@ -1440,11 +1640,12 @@ def circuit(x, y): def test_autograph_jacobian(self, decorator, interface, tol): """Test that a parameter-shift vector-valued QNode can be compiled using @tf.function, and differentiated""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1482,14 +1683,10 @@ def circuit(x, y): def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) def circuit(x): qml.RY(x, wires=0) return qml.expval(qml.PauliZ(0)) @@ -1510,14 +1707,10 @@ def circuit(x): def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -1536,44 +1729,16 @@ def circuit(x): expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)] assert np.allclose(g, expected_g, atol=tol, rtol=0) - @pytest.mark.xfail - @pytest.mark.parametrize("grad_on_execution", [True, False]) - def test_autograph_adjoint_multi_out(self, grad_on_execution, decorator, interface, tol): - """Test that a parameter-shift QNode can be compiled - using @tf.function, and differentiated to second order""" - - @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) - - x = tf.Variable(0.5, dtype=tf.float64) - - with tf.GradientTape() as tape: - res = qml.math.hstack(circuit(x)) - g = tape.jacobian(res, x) - a = x * 1.0 - - expected_res = [tf.cos(a), tf.sin(a)] - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-tf.sin(a), tf.cos(a)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - def test_autograph_ragged_differentiation(self, decorator, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1605,11 +1770,12 @@ def circuit(x, y): def test_autograph_hessian(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) a = tf.Variable(0.543, dtype=tf.float64) b = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=2, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=2, interface=interface) def circuit(x, y): qml.RY(x, wires=0) qml.RX(y, wires=0) @@ -1635,15 +1801,35 @@ def circuit(x, y): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) + def test_autograph_sample(self, decorator, interface): + """Test that a QNode returning raw samples can be compiled + using @tf.function""" + dev = qml.device("default.qubit", wires=2, shots=100) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + @decorator + @qnode(dev, diff_method="parameter-shift", interface=interface) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.sample() + + result = circuit(x, y) + result = np.array(result).flatten() + assert np.all([r in [1, 0] for r in result]) + def test_autograph_state(self, decorator, interface, tol): """Test that a parameter-shift QNode returning a state can be compiled using @tf.function""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) # TODO: fix this for diff_method=None @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + @qnode(dev, diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1660,15 +1846,16 @@ def circuit(x, y): def test_autograph_dimension(self, decorator, interface): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) - def circuit(**_): + @qnode(dev, diff_method="parameter-shift", interface=interface) + def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) # pylint:disable=unexpected-keyword-arg + res = circuit() res = tf.stack(res) assert res.shape == (2, 10) @@ -1676,23 +1863,24 @@ def circuit(**_): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and one param, the gradient is a float.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1710,16 +1898,19 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1740,16 +1931,19 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1768,64 +1962,68 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) + a = tf.Variable(0.1) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) - jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, (a, b)) assert isinstance(jac, tuple) @@ -1836,94 +2034,101 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single array is returned.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) + a = tf.Variable([0.1, 0.2]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a single params return an array.""" + num_wires = 2 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) + a = tf.Variable(0.1) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1935,40 +2140,50 @@ def circuit(a, b): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) + a = tf.Variable([0.1, 0.2]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (5, 2) def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -1981,7 +2196,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2008,13 +2222,20 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2023,7 +2244,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2042,10 +2262,10 @@ def circuit(x): assert isinstance(hess, tf.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, interface): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2061,7 +2281,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2088,7 +2307,7 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2097,6 +2316,8 @@ def test_hessian_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") + dev = qml.device(dev_name, wires=2) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2105,7 +2326,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2125,9 +2345,16 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + pytest.skip("Test does not support hadamard because multiple measurements.") + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2140,7 +2367,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2168,7 +2394,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" @@ -2178,6 +2404,10 @@ def test_hessian_probs_expval_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because multiple measurements.") + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2186,7 +2416,6 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2207,9 +2436,11 @@ def circuit(x): assert hess.shape == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") if diff_method == "hadamard": @@ -2224,7 +2455,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2252,7 +2482,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2260,6 +2490,8 @@ def test_hessian_probs_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") + dev = qml.device(dev_name, wires=2) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2268,7 +2500,6 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2289,49 +2520,17 @@ def circuit(x): assert hess.shape == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="tf") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="tf") def circuit(): qml.Hadamard(wires=0) return qml.state() res = circuit() assert isinstance(res, tf.Tensor) - - -def test_error_device_vjp_jacobian(): - """Test a ValueError is raised if a jacobian is attempted to be computed with device_vjp=True.""" - - dev = qml.device("default.qubit") - - @qml.qnode(dev, diff_method="adjoint", device_vjp=True) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) - - x = tf.Variable(0.1) - - with tf.GradientTape() as tape: - y = qml.math.hstack(circuit(x)) - - with pytest.raises(ValueError): - tape.jacobian(y, x) - - -def test_error_device_vjp_state_float32(): - """Test a ValueError is raised is state differentiation is attemped with float32 parameters.""" - - dev = qml.device("default.qubit") - - @qml.qnode(dev, diff_method="adjoint", device_vjp=True) - def circuit(x): - qml.RX(x, wires=0) - return qml.probs(wires=0) - - x = tf.Variable(0.1, dtype=tf.float32) - with pytest.raises(ValueError, match="tensorflow with adjoint differentiation of the state"): - with tf.GradientTape(): - circuit(x) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py similarity index 82% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py index a3a5606fc78..bea3056b7c5 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,unexpected-keyword-arg +# pylint: disable=too-many-arguments import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -28,9 +27,9 @@ shots_and_num_copies_hess = [((10, (5, 1)), 2)] qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", {"h": 10e-2}], - [DefaultQubit(), "parameter-shift", {}], - [DefaultQubit(), "spsa", {"h": 10e-2, "num_directions": 20}], + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], ] TOLS = { @@ -46,15 +45,16 @@ @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -65,7 +65,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,9 +74,10 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -88,7 +89,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -100,9 +101,10 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -113,7 +115,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -122,10 +124,11 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -136,7 +139,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -145,10 +148,11 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -160,7 +164,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -172,10 +176,11 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -186,7 +191,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -195,9 +200,10 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1) par_1 = tf.Variable(0.2) @@ -210,7 +216,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -222,9 +228,10 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -236,7 +243,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2, 0.3]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -245,9 +252,10 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -258,7 +266,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -267,9 +275,10 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -281,7 +290,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -293,9 +302,10 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -306,7 +316,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -318,15 +328,16 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -340,7 +351,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -355,9 +366,10 @@ def circuit(x, y): assert h.shape == (2, num_copies) def test_hessian_expval_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) params = tf.Variable([0.1, 0.2], dtype=tf.float64) @@ -370,7 +382,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params, shots=shots) + res = circuit(params) res = qml.math.stack(res) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -381,9 +393,10 @@ def circuit(x): assert hess.shape == (num_copies, 2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -397,7 +410,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -412,10 +425,12 @@ def circuit(x, y): assert h.shape == (2, num_copies, 3) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) @@ -427,7 +442,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params, shots=shots) + res = circuit(params) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -438,22 +453,22 @@ def circuit(x): assert hess.shape == (num_copies, 3, 2, 2) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((20000, 18000, 16000), 3), ((20000, (18000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -465,7 +480,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -482,11 +497,11 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -498,7 +513,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/legacy_devices_integration/test_torch_legacy.py b/tests/interfaces/legacy_devices_integration/test_torch_legacy.py new file mode 100644 index 00000000000..5aac97fb0e3 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_torch_legacy.py @@ -0,0 +1,1308 @@ +# Copyright 2018-2022 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the Torch interface""" +# pylint: disable=protected-access,too-few-public-methods +import numpy as np +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import finite_diff, param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.torch +torch = pytest.importorskip("torch") +torch_functional = pytest.importorskip("torch.autograd.functional") +torch_cuda = pytest.importorskip("torch.cuda") + + +@pytest.mark.parametrize("interface", ["torch", "auto"]) +class TestTorchExecuteUnitTests: + """Unit tests for torch execution""" + + def test_jacobian_options(self, interface, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute( + [tape], + dev, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + interface=interface, + )[0] + + res.backward() + + for args in spy.call_args_list: + assert args[1]["shift"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self, interface): + """Test that an error is raised if a gradient transform + is used with grad_on_execution=True""" + a = torch.tensor([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + execute( + [tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface=interface + ) + + def test_grad_on_execution_reuse_state(self, interface, mocker): + """Test that grad_on_execution uses the `device.execute_and_gradients` pathway + while reusing the quantum state.""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface=interface, + ) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_grad_on_execution(self, interface, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian"}, + interface=interface, + ) + + # two device executions; one for the value, one for the Jacobian + assert dev.num_executions == 2 + spy.assert_called() + + def test_no_grad_on_execution(self, interface, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface=interface, + )[0] + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + res.backward() + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], dev, gradient_fn=param_shift, cachesize=cachesize, interface="torch" + )[0][0] + + params = torch.tensor([0.1, 0.2], requires_grad=True) + res = cost(params, cachesize=2) + res.backward() + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ + 0 + ] + + custom_cache = {} + params = torch.tensor([0.1, 0.2], requires_grad=True) + res = cost(params, cache=custom_cache) + res.backward() + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self): + """Test that, with the parameter-shift transform, + Torch always uses the optimum number of evals when computing the Jacobian.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ + 0 + ] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + params = torch.tensor([0.1, 0.2], requires_grad=True) + torch_functional.jacobian(lambda p: cost(p, cache=None), params) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + torch_functional.jacobian(lambda p: cost(p, cache=True), params) + assert dev.num_executions == 5 + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, with the parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor(np.arange(1, num_params + 1) / 10, requires_grad=True) + + N = len(params) + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], dev, gradient_fn=param_shift, cache=cache, interface="torch", max_diff=2 + )[0] + + # No caching: number of executions is not ideal + hess1 = torch.autograd.functional.hessian(lambda x: cost(x, cache=None), params) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params.detach() + expected = torch.tensor( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + expected_runs = 1 # forward pass + expected_runs += 2 * N # Jacobian + expected_runs += 4 * N + 1 # Hessian diagonal + expected_runs += 4 * N**2 # Hessian off-diagonal + assert dev.num_executions == expected_runs + + # Use caching: number of executions is ideal + dev._num_executions = 0 + hess2 = torch.autograd.functional.hessian(lambda x: cost(x, cache=True), params) + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + def test_caching_adjoint_no_grad_on_execution(self): + """Test that caching reduces the number of adjoint evaluations + when grad_on_execution=False""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface="torch", + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one for the backward pass + torch_functional.jacobian(lambda x: cost(x, cache=None), params) + assert dev.num_executions == 2 + + # With caching, only 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + torch_functional.jacobian(lambda x: cost(x, cache=True), params) + assert dev.num_executions == 2 + + +torch_devices = [None] + +if torch_cuda.is_available(): + torch_devices.append(torch.device("cuda")) + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift, "interface": "torch"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, + "interface": "torch", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "torch", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "torch", + }, + {"gradient_fn": param_shift, "interface": "auto"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "auto", + }, +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("torch_device", torch_devices) +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestTorchExecuteIntegration: + """Test the torch interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, torch_device, execute_kwargs): + """Test that the execute function produces results with the expected shapes""" + dev = qml.device("default.qubit.legacy", wires=1) + a = torch.tensor(0.1, requires_grad=True, device=torch_device) + b = torch.tensor(0.2, requires_grad=False, device=torch_device) + + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + res = execute([tape1, tape2], dev, **execute_kwargs) + + assert isinstance(res, TensorLike) + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, torch_device, execute_kwargs, tol): + """Test scalar jacobian calculation by comparing two types of pipelines""" + a = torch.tensor(0.1, requires_grad=True, dtype=torch.float64, device=torch_device) + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + res.backward() + + # compare to backprop gradient + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + dev = qml.device("default.qubit.autograd", wires=2) + return dev.batch_execute([tape])[0] + + expected = qml.grad(cost, argnum=0)(0.1) + assert torch.allclose(a.grad, torch.tensor(expected, device=torch_device), atol=tol, rtol=0) + + def test_jacobian(self, torch_device, execute_kwargs, tol): + """Test jacobian calculation by checking against analytic values""" + a_val = 0.1 + b_val = 0.2 + + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + b = torch.tensor(b_val, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RZ(torch.tensor(0.543, device=torch_device), wires=0) + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1, 2] + + assert isinstance(res, tuple) + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == () + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == () + + expected = torch.tensor( + [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)], device=torch_device + ) + assert torch.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert torch.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + loss = res[0] + res[1] + + loss.backward() + expected = torch.tensor( + [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)], + dtype=a.dtype, + device=torch_device, + ) + assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) + assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + + def test_tape_no_parameters(self, torch_device, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + dev = qml.device("default.qubit.legacy", wires=1) + params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device) + x, y = params.detach() + + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(0.5, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + + res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) + expected = torch.tensor(1 + np.cos(0.5), dtype=res.dtype) + torch.cos(x) * torch.cos(y) + expected = expected.to(device=res.device) + + assert torch.allclose(res, expected, atol=tol, rtol=0) + + res.backward() + grad = params.grad.detach() + expected = torch.tensor( + [-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)], + dtype=grad.dtype, + device=grad.device, + ) + assert torch.allclose(grad, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = torch.tensor(0.1, requires_grad=True, device=torch_device) + b = torch.tensor(0.2, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + assert tape.trainable_params == [0, 1] + + res = execute([tape], dev, **execute_kwargs)[0] + loss = res[0] + res[1] + loss.backward() + + a_val = 0.54 + b_val = 0.8 + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + b = torch.tensor(b_val, requires_grad=True, device=torch_device) + + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + + expected = torch.tensor( + [np.cos(2 * a_val), -np.cos(2 * a_val) * np.sin(b_val)], + device=torch_device, + dtype=res2[0].dtype, + ) + assert torch.allclose(res2[0].detach(), expected[0], atol=tol, rtol=0) + assert torch.allclose(res2[1].detach(), expected[1], atol=tol, rtol=0) + + loss = res2[0] + res2[1] + loss.backward() + + expected = torch.tensor( + [ + -2 * np.sin(2 * a_val) + 2 * np.sin(2 * a_val) * np.sin(b_val), + -np.cos(2 * a_val) * np.cos(b_val), + ], + dtype=a.dtype, + device=torch_device, + ) + + assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) + assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + + def test_classical_processing(self, torch_device, execute_kwargs, tol): + """Test the classical processing of gate parameters within the quantum tape""" + p_val = [0.1, 0.2] + params = torch.tensor(p_val, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(params[0] * params[1], wires=0) + qml.RZ(0.2, wires=0) + qml.RX(params[1] + params[1] ** 2 + torch.sin(params[0]), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert tape.trainable_params == [0, 2] + + tape_params = torch.tensor([i.detach() for i in tape.get_parameters()], device=torch_device) + expected = torch.tensor( + [p_val[0] * p_val[1], p_val[1] + p_val[1] ** 2 + np.sin(p_val[0])], + dtype=tape_params.dtype, + device=torch_device, + ) + + assert torch.allclose( + tape_params, + expected, + atol=tol, + rtol=0, + ) + + res.backward() + + assert isinstance(params.grad, torch.Tensor) + assert params.shape == (2,) + + def test_no_trainable_parameters(self, torch_device, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(0.2, wires=0) + qml.RX(torch.tensor(0.1, device=torch_device), wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [] + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == () + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == () + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res[0].backward() + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res[1].backward() + + @pytest.mark.parametrize( + "U", [torch.tensor([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]])] + ) + def test_matrix_parameter(self, torch_device, U, execute_kwargs, tol): + """Test that the torch interface works correctly + with a matrix parameter""" + a_val = 0.1 + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + + if isinstance(U, torch.Tensor) and torch_device is not None: + U = U.to(torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1] + + expected = torch.tensor(-np.cos(a_val), dtype=res.dtype, device=torch_device) + assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor([np.sin(a_val)], dtype=a.grad.dtype, device=torch_device) + assert torch.allclose(a.grad, expected, atol=tol, rtol=0) + + def test_differentiable_expand(self, torch_device, execute_kwargs, tol): + """Test that operation and nested tape expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + dev = qml.device("default.qubit.legacy", wires=1) + a = np.array(0.1) + p_val = [0.1, 0.2, 0.3] + p = torch.tensor(p_val, requires_grad=True, device=torch_device) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(a, wires=0) + U3(p[0], p[1], p[2], wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + expected = torch.tensor( + np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) + + np.sin(a) + * ( + np.cos(p_val[2]) * np.sin(p_val[1]) + + np.cos(p_val[0]) * np.cos(p_val[1]) * np.sin(p_val[2]) + ), + dtype=res.dtype, + device=torch_device, + ) + assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [ + np.cos(p_val[1]) + * (np.cos(a) * np.cos(p_val[0]) - np.sin(a) * np.sin(p_val[0]) * np.sin(p_val[2])), + np.cos(p_val[1]) * np.cos(p_val[2]) * np.sin(a) + - np.sin(p_val[1]) + * (np.cos(a) * np.sin(p_val[0]) + np.cos(p_val[0]) * np.sin(a) * np.sin(p_val[2])), + np.sin(a) + * ( + np.cos(p_val[0]) * np.cos(p_val[1]) * np.cos(p_val[2]) + - np.sin(p_val[1]) * np.sin(p_val[2]) + ), + ], + dtype=p.grad.dtype, + device=torch_device, + ) + assert torch.allclose(p.grad, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, torch_device, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True, device=torch_device) + y = torch.tensor(y_val, requires_grad=True, device=torch_device) + + def circuit(x, y): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = circuit(x, y) + + expected_0 = torch.tensor( + [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2], + dtype=res[0].dtype, + device=torch_device, + ) + + expected_1 = torch.tensor( + [ + (1 + np.cos(x_val) * np.cos(y_val)) / 2, + (1 - np.cos(x_val) * np.cos(y_val)) / 2, + ], + dtype=res[0].dtype, + device=torch_device, + ) + + assert torch.allclose(res[0], expected_0, atol=tol, rtol=0) + assert torch.allclose(res[1], expected_1, atol=tol, rtol=0) + + jac = torch_functional.jacobian(circuit, (x, y)) + dtype_jac = jac[0][0].dtype + + res_0 = torch.tensor( + [-np.sin(x_val) / 2, np.sin(x_val) / 2], dtype=dtype_jac, device=torch_device + ) + res_1 = torch.tensor([0.0, 0.0], dtype=dtype_jac, device=torch_device) + res_2 = torch.tensor( + [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + res_3 = torch.tensor( + [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + + assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) + assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) + assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) + assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + + def test_ragged_differentiation(self, torch_device, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True, device=torch_device) + y = torch.tensor(y_val, requires_grad=True, device=torch_device) + + def circuit(x, y): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = circuit(x, y) + + res_0 = torch.tensor(np.cos(x_val), dtype=res[0].dtype, device=torch_device) + res_1 = torch.tensor( + [(1 + np.cos(x_val) * np.cos(y_val)) / 2, (1 - np.cos(x_val) * np.cos(y_val)) / 2], + dtype=res[0].dtype, + device=torch_device, + ) + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert torch.allclose(res[0], res_0, atol=tol, rtol=0) + assert torch.allclose(res[1], res_1, atol=tol, rtol=0) + + jac = torch_functional.jacobian(circuit, (x, y)) + dtype_jac = jac[0][0].dtype + + res_0 = torch.tensor( + -np.sin(x_val), + dtype=dtype_jac, + device=torch_device, + ) + res_1 = torch.tensor( + 0.0, + dtype=dtype_jac, + device=torch_device, + ) + res_2 = torch.tensor( + [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + res_3 = torch.tensor( + [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + + assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) + assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) + assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) + assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + + def test_sampling(self, torch_device, execute_kwargs): + """Test sampling works as expected""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + if execute_kwargs["interface"] == "auto": + pytest.skip("Can't detect interface without a parametrized gate in the tape") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == (10,) + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == (10,) + + def test_sampling_expval(self, torch_device, execute_kwargs): + """Test sampling works as expected if combined with expectation values""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + if execute_kwargs["interface"] == "auto": + pytest.skip("Can't detect interface without a parametrized gate in the tape") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.expval(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert len(res) == 2 + assert isinstance(res, tuple) + assert res[0].shape == (10,) + assert res[1].shape == () + assert isinstance(res[0], torch.Tensor) + assert isinstance(res[1], torch.Tensor) + + def test_sampling_gradient_error(self, torch_device, execute_kwargs): + """Test differentiating a tape with sampling results in an error""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + dev = qml.device("default.qubit.legacy", wires=1, shots=10) + + x = torch.tensor(0.65, requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.sample() + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res.backward() + + def test_repeated_application_after_expand(self, torch_device, execute_kwargs): + """Test that the Torch interface continues to work after + tape expansions""" + # pylint: disable=unused-argument + n_qubits = 2 + dev = qml.device("default.qubit.legacy", wires=n_qubits) + + weights = torch.ones((3,)) + + with qml.queuing.AnnotatedQueue() as q: + qml.U3(*weights, wires=0) + qml.expval(qml.PauliZ(wires=0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape = tape.expand() + execute([tape], dev, **execute_kwargs) + + +@pytest.mark.parametrize("torch_device", torch_devices) +class TestHigherOrderDerivatives: + """Test that the torch execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + torch.tensor([0.543, -0.654], requires_grad=True), + torch.tensor([0, -0.654], requires_grad=True), + torch.tensor([-2.0, 0], requires_grad=True), + ], + ) + def test_parameter_shift_hessian(self, torch_device, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using torch, yielding second derivatives.""" + # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, interface="torch", max_diff=2 + ) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params.detach() + expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert torch.allclose(params.grad.detach(), expected, atol=tol, rtol=0) + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.tensor( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + def test_hessian_vector_valued(self, torch_device, tol): + """Test hessian calculation of a vector valued tape""" + dev = qml.device("default.qubit.legacy", wires=1) + + def circuit(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(x[0], wires=0) + qml.RX(x[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return torch.stack( + execute([tape], dev, gradient_fn=param_shift, interface="torch", max_diff=2) + ) + + x = torch.tensor([1.0, 2.0], requires_grad=True, device=torch_device) + res = circuit(x) + + if torch_device is not None: + a, b = x.detach().cpu().numpy() + else: + a, b = x.detach().numpy() + + expected_res = torch.tensor( + [ + 0.5 + 0.5 * np.cos(a) * np.cos(b), + 0.5 - 0.5 * np.cos(a) * np.cos(b), + ], + dtype=res.dtype, + device=torch_device, + ) + assert torch.allclose(res.detach(), expected_res, atol=tol, rtol=0) + + def jac_fn(x): + return torch_functional.jacobian(circuit, x, create_graph=True) + + g = jac_fn(x) + + hess = torch_functional.jacobian(jac_fn, x) + + expected_g = torch.tensor( + [ + [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)], + [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)], + ], + dtype=g.dtype, + device=torch_device, + ) + assert torch.allclose(g.detach(), expected_g, atol=tol, rtol=0) + + expected_hess = torch.tensor( + [ + [ + [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)], + [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)], + ], + [ + [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)], + [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)], + ], + ], + dtype=hess.dtype, + device=torch_device, + ) + assert torch.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) + + def test_adjoint_hessian(self, torch_device, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor( + [0.543, -0.654], requires_grad=True, dtype=torch.float64, device=torch_device + ) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface="torch", + )[0] + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.zeros([2, 2], dtype=torch.float64, device=torch_device) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + def test_max_diff(self, torch_device, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface="torch" + ) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params.detach() + expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert torch.allclose( + params.grad.detach().to(torch_device), expected.to(torch_device), atol=tol, rtol=0 + ) + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.zeros([2, 2], dtype=torch.float64) + assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift, "interface": "torch"}, + {"gradient_fn": finite_diff, "interface": "torch"}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + + return torch.hstack(execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1.detach().numpy() + d = coeffs2.detach().numpy()[0] + x, y = weights.detach().numpy() + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1.detach().numpy() + d = coeffs2.detach().numpy()[0] + x, y = weights.detach().numpy() + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False, dtype=torch.float64) + coeffs2 = torch.tensor([0.7], requires_grad=False, dtype=torch.float64) + weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + res = torch.hstack( + torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) + ) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64) + coeffs2 = torch.tensor([0.7], requires_grad=True, dtype=torch.float64) + weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + res = torch.hstack( + torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) + ) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res.detach(), expected, atol=tol, rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py similarity index 74% rename from tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py index 2f1c85f6275..3965255a4a4 100644 --- a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,15 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the Torch interface with a QNode""" -# pylint: disable=too-many-arguments,unexpected-keyword-arg,no-member,comparison-with-callable, no-name-in-module -# pylint: disable=use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment, use-dict-literal +# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods, +# pylint: disable=use-dict-literal, use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment import numpy as np import pytest import pennylane as qml from pennylane import qnode -from pennylane.devices import DefaultQubit -from tests.param_shift_dev import ParamShiftDerivativesDevice pytestmark = pytest.mark.torch @@ -28,25 +26,14 @@ jacobian = torch.autograd.functional.jacobian hessian = torch.autograd.functional.hessian -# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", True, True], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, False], - [ParamShiftDerivativesDevice(), "best", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_and_qubit_device_and_diff_method = [ @@ -59,25 +46,25 @@ @pytest.mark.parametrize( - "interface, dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface, dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with Torch integrates with the PennyLane stack""" - def test_execution_with_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a): qml.RY(a, wires=0) @@ -92,7 +79,7 @@ def circuit(a): # with the interface, the tape returns torch tensors assert isinstance(res, torch.Tensor) - assert res.shape == () + assert res.shape == tuple() # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -102,18 +89,20 @@ def circuit(a): grad = a.grad assert isinstance(grad, torch.Tensor) - assert grad.shape == () + assert grad.shape == tuple() - def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_interface_swap(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the Torch interface can be applied to a QNode with a pre-existing interface""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - interface="autograd", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface="autograd", grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -139,19 +128,22 @@ def circuit(a): assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_drawing(self, interface, dev_name, diff_method, grad_on_execution): """Test circuit drawing when using the torch interface""" x = torch.tensor(0.1, requires_grad=True) y = torch.tensor([0.2, 0.3], requires_grad=True) z = torch.tensor(0.4, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) @@ -167,17 +159,14 @@ def circuit(p1, p2=y, **kwargs): assert result == expected - def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation""" kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a_val = 0.1 @@ -186,6 +175,13 @@ def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_v a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -221,29 +217,19 @@ def circuit(a, b): assert np.allclose(b.grad, expected[1], atol=tol, rtol=0) # TODO: fix this behavior with float: already present before return type. - def test_jacobian_dtype( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + @pytest.mark.xfail + def test_jacobian_dtype(self, interface, dev_name, diff_method, grad_on_execution): """Test calculating the jacobian with a different datatype""" - if not "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Failing unless lightning.qubit") if diff_method == "backprop": pytest.skip("Test does not support backprop") a = torch.tensor(0.1, dtype=torch.float32, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True) + dev = qml.device(dev_name, wires=2) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -267,20 +253,15 @@ def circuit(a, b): assert a.grad.dtype is torch.float32 assert b.grad.dtype is torch.float32 - def test_jacobian_options( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): """Test setting jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite-diff and spsa") a = torch.tensor([0.1, 0.2], requires_grad=True) + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, @@ -288,7 +269,6 @@ def test_jacobian_options( interface=interface, h=1e-8, approx_order=2, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -298,9 +278,7 @@ def circuit(a): res = circuit(a) res.backward() - def test_changing_trainability( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -312,12 +290,10 @@ def test_changing_trainability( a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) + dev = qml.device(dev_name, wires=2) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -366,25 +342,21 @@ def circuit(a, b): expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val) assert np.allclose(a.grad, expected, atol=tol, rtol=0) - def test_classical_processing( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): """Test classical processing within the quantum tape""" a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False) c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -404,22 +376,17 @@ def circuit(a, b, c): assert b.grad is None assert isinstance(c.grad, torch.Tensor) - def test_no_trainable_parameters( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): """Test evaluation and Jacobian if there are no trainable parameters""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b): qml.RY(a, wires=0) @@ -463,20 +430,21 @@ def circuit(a, b): np.array([[0, 1], [1, 0]]), ], ) - def test_matrix_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, U, tol - ): + def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, U, tol): """Test that the Torch interface works correctly with a matrix parameter""" a_val = 0.1 a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -493,23 +461,18 @@ def circuit(U, a): res.backward() assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0) - def test_differentiable_expand( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that operation and nested tapes expansion is differentiable""" kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - class U3(qml.U3): # pylint:disable=too-few-public-methods + class U3(qml.U3): def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -518,6 +481,12 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) a = np.array(0.1) p_val = [0.1, 0.2, 0.3] p = torch.tensor(p_val, dtype=torch.float64, requires_grad=True) @@ -558,9 +527,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self): + def test_changing_shots(self, mocker, tol): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -568,22 +537,33 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # only same call as above # TODO: add this test after shot vectors addition @pytest.mark.xfail def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = DefaultQubit() + # pylint: disable=unexpected-keyword-arg + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = torch.tensor([0.543, -0.654], requires_grad=True) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -605,10 +585,11 @@ def test_multiple_gradient_integration(self, tol): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) weights = torch.tensor([0.543, -0.654], requires_grad=True) a, b = weights - @qnode(DefaultQubit(), interface="torch", diff_method=qml.gradients.param_shift) + @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -616,7 +597,7 @@ def circuit(a, b): return qml.expval(qml.PauliY(1)) res1 = circuit(*weights) - assert qml.math.shape(res1) == () + assert qml.math.shape(res1) == tuple() res2 = circuit(*weights, shots=[(1, 1000)]) assert len(res2) == 1000 @@ -628,47 +609,120 @@ def circuit(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = torch.tensor([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface="torch") + @qnode(dev, interface="torch") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used cost_fn(a, b) + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + # original QNode settings are unaffected + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + +class TestAdjoint: + """Specific integration tests for the adjoint method""" + + def test_reuse_state(self, mocker): + """Tests that the Torch interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, diff_method="adjoint", interface="torch") + def circ(x): + qml.RX(x[0], wires=0) + qml.RY(x[1], wires=1) + qml.CNOT(wires=(0, 1)) + return qml.expval(qml.PauliZ(0)) + + expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])]) + + spy = mocker.spy(dev, "adjoint_jacobian") + + x1 = torch.tensor([1.0, 1.0], requires_grad=True) + res = circ(x1) + res.backward() + + assert np.allclose(x1.grad[0], expected_grad(x1)[0]) + assert circ.device.num_executions == 1 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + + def test_resuse_state_multiple_evals(self, mocker, tol): + """Tests that the Torch interface reuses the device state for adjoint differentiation, + even where there are intermediate evaluations.""" + dev = qml.device("default.qubit.legacy", wires=2) + + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True) + y = torch.tensor(y_val, requires_grad=True) + + @qnode(dev, diff_method="adjoint", interface="torch") + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + res1 = circuit(x, y) + assert np.allclose(res1.detach(), np.cos(x_val), atol=tol, rtol=0) + + # intermediate evaluation with different values + circuit(torch.tan(x), torch.cosh(y)) + + # the adjoint method will continue to compute the correct derivative + res1.backward() + assert np.allclose(x.grad.detach(), -np.sin(x_val), atol=tol, rtol=0) + assert dev.num_executions == 2 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = {} - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -679,7 +733,6 @@ def test_probability_differentiation( diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -715,24 +768,25 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_ragged_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -770,22 +824,17 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_chained_qnodes( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit1(weights): qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) @@ -817,11 +866,18 @@ def cost(weights): loss = cost(weights) loss.backward() - def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a scalar valued QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -832,7 +888,6 @@ def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vj grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -868,13 +923,18 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a vector valued QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -885,7 +945,6 @@ def test_hessian_vector_valued( grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -932,11 +991,18 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a ragged QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -947,7 +1013,6 @@ def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, de grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -1006,23 +1071,29 @@ def circuit_stack(x): assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") options = {} if diff_method == "finite-diff": options = {"h": 1e-6} + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -1067,23 +1138,23 @@ def cost_fn(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 - if dev.name == "param_shift.qubit": - pytest.skip("parameter shift does not support measuring the state.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning devices do not support state with adjoint diff.") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = torch.tensor(0.543, requires_grad=True) y = torch.tensor(-0.654, requires_grad=True) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1093,7 +1164,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert torch.is_complex(res) + assert res.dtype is torch.complex128 probs = torch.abs(res) ** 2 return probs[0] + probs[2] @@ -1110,25 +1181,20 @@ def cost_fn(x, y): assert torch.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the variance of a projector is correctly returned""" kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements") - if diff_method == "spsa": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard does not support variances.") + dev = qml.device(dev_name, wires=2) P = torch.tensor(state, requires_grad=False) x, y = 0.765, -0.654 @@ -1156,64 +1222,99 @@ def circuit(x, y): ) assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = torch.tensor(0.543, dtype=torch.float64, requires_grad=True) + phi = torch.tensor(-0.654, dtype=torch.float64, requires_grad=True) + + @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = torch.exp(2 * r) * torch.sin(phi) ** 2 + torch.exp(-2 * r) * torch.cos(phi) ** 2 + assert torch.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res.backward() + res = torch.tensor([r.grad, phi.grad]) + expected = torch.tensor( + [ + [ + 2 * torch.exp(2 * r) * torch.sin(phi) ** 2 + - 2 * torch.exp(-2 * r) * torch.cos(phi) ** 2, + 2 * torch.sinh(2 * r) * torch.sin(2 * phi), + ] + ] ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + assert torch.allclose(res, expected, atol=tol, rtol=0) - phi = torch.tensor(1.23, requires_grad=True) - theta = torch.tensor(4.56, requires_grad=True) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA - assert qml.math.allclose(circuit(phi, theta), expected_circuit(theta)) + n = torch.tensor(0.12, dtype=torch.float64, requires_grad=True) + a = torch.tensor(0.765, dtype=torch.float64, requires_grad=True) - gradient = torch.autograd.grad(circuit(phi, theta), [phi, theta]) - exp_theta_grad = torch.autograd.grad(expected_circuit(theta), theta)[0] - assert qml.math.allclose(gradient, [0.0, exp_theta_grad]) + @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + res = circuit(n, a) + expected = n**2 + n + torch.abs(a) ** 2 * (1 + 2 * n) + assert torch.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method -) + # circuit jacobians + res.backward() + res = torch.tensor([n.grad, a.grad]) + expected = torch.tensor([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Torch interface""" - def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp): + def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1224,7 +1325,6 @@ def decomposition(self): diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, interface="torch", ) def circuit(x): @@ -1235,32 +1335,33 @@ def circuit(x): x = torch.tensor(0.5, requires_grad=True, dtype=torch.float64) loss = circuit(x) - loss.backward() res = x.grad assert torch.allclose(res, -3 * torch.sin(3 * x)) - if diff_method == "parameter-shift" and dev.name != "param_shift.qubit": + if diff_method == "parameter-shift": # test second order derivatives res = torch.autograd.functional.hessian(circuit, x) assert torch.allclose(res, -9 * torch.cos(3 * x)) @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, - dev, - diff_method, - grad_on_execution, - device_vjp, - max_diff, + self, dev_name, diff_method, grad_on_execution, max_diff ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1272,7 +1373,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1288,7 +1388,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp, tol + self, dev_name, diff_method, grad_on_execution, max_diff, tol ): """Test that if there are non-commuting groups and the number of shots is None @@ -1298,17 +1398,17 @@ def test_hamiltonian_expansion_analytic( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, ) if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1348,11 +1448,7 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if ( - diff_method in ("parameter-shift", "backprop") - and max_diff == 2 - and dev.name != "param_shift.qubit" - ): + if diff_method in ("parameter-shift", "backprop") and max_diff == 2: hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] @@ -1372,14 +1468,14 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, device_vjp, grad_on_execution, max_diff + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1391,10 +1487,11 @@ def test_hamiltonian_finite_shots( tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - tol = 0.15 elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1403,7 +1500,6 @@ def test_hamiltonian_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1418,17 +1514,15 @@ def circuit(data, weights, coeffs): c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin( d[1] + w[1] ) assert torch.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - res.backward() grad = (w.grad, c.grad) @@ -1447,10 +1541,8 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": - hessians = torch.autograd.functional.hessian( - lambda _d, _w, _c: circuit(_d, _w, _c, shots=50000), (d, w, c) - ) + if diff_method == "parameter-shift" and max_diff == 2: + hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64), atol=tol) @@ -1474,14 +1566,15 @@ class TestSample: def test_sample_dimension(self): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 @@ -1494,33 +1587,37 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (10,) + assert res[0].shape == (shots,) assert isinstance(res[1], torch.Tensor) assert res[1].shape == () def test_counts_expval(self): """Test counts works as expected if combined with expectation values""" + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(qml.device("default.qubit"), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) @@ -1531,19 +1628,22 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, tuple) assert len(result) == 3 - assert np.array_equal(result[0].shape, (10,)) + assert np.array_equal(result[0].shape, (n_sample,)) assert result[1].shape == () assert isinstance(result[1], torch.Tensor) assert result[2].shape == () @@ -1552,73 +1652,73 @@ def circuit(): def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" + n_sample = 10 - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) + + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, torch.Tensor) - assert np.array_equal(result.shape, (10,)) + assert np.array_equal(result.shape, (n_sample,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=10) + result = circuit() # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) - assert tuple(result[0].shape) == (10,) - assert tuple(result[1].shape) == (10,) - assert tuple(result[2].shape) == (10,) + assert tuple(result[0].shape) == (n_sample,) + assert tuple(result[1].shape) == (n_sample,) + assert tuple(result[2].shape) == (n_sample,) assert result[0].dtype == torch.float64 assert result[1].dtype == torch.float64 assert result[2].dtype == torch.float64 qubit_device_and_diff_method_and_grad_on_execution = [ - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", True, True], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "hadamard", False, False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "hadamard", False], ] @pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", - qubit_device_and_diff_method_and_grad_on_execution, + "dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method_and_grad_on_execution ) @pytest.mark.parametrize("shots", [None, 10000]) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - # pylint:disable=too-many-public-methods - - def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots - ): + def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution, shots): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1626,7 +1726,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - res = circuit(a, shots=shots) + res = circuit(a) assert isinstance(res, torch.Tensor) assert res.shape == () @@ -1638,19 +1738,20 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, shots, device_vjp + self, dev_name, diff_method, grad_on_execution, shots ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1659,7 +1760,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - res = circuit(a, b, shots=shots) + res = circuit(a, b) # gradient res.backward() @@ -1670,19 +1771,20 @@ def circuit(a, b): assert grad_b.shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1690,7 +1792,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - res = circuit(a, shots=shots) + res = circuit(a) # gradient res.backward() @@ -1700,7 +1802,7 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1710,13 +1812,14 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1730,7 +1833,7 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1739,13 +1842,14 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1754,7 +1858,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) + jac = jacobian(circuit, (a, b)) assert isinstance(jac, tuple) @@ -1765,7 +1869,7 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1774,31 +1878,39 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, torch.Tensor) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1807,7 +1919,6 @@ def test_jacobian_expval_expval_multiple_params( interface="torch", diff_method=diff_method, max_diff=1, - device_vjp=device_vjp, grad_on_execution=grad_on_execution, ) def circuit(x, y): @@ -1816,7 +1927,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + jac = jacobian(circuit, (par_0, par_1)) assert isinstance(jac, tuple) @@ -1835,19 +1946,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1855,7 +1967,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1867,7 +1979,7 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, device_vjp, grad_on_execution, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1877,6 +1989,8 @@ def test_jacobian_var_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1884,7 +1998,6 @@ def test_jacobian_var_var_multiple_params( dev, interface="torch", diff_method=diff_method, - device_vjp=device_vjp, max_diff=1, grad_on_execution=grad_on_execution, ) @@ -1894,7 +2007,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + jac = jacobian(circuit, (par_0, par_1)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1914,7 +2027,7 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, device_vjp, grad_on_execution, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -1924,13 +2037,9 @@ def test_jacobian_var_var_multiple_params_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1938,7 +2047,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1950,22 +2059,22 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1973,7 +2082,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1985,7 +2094,7 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1993,13 +2102,14 @@ def test_jacobian_multiple_measurement_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2008,7 +2118,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) + jac = jacobian(circuit, (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2028,7 +2138,7 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2036,13 +2146,14 @@ def test_jacobian_multiple_measurement_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2050,7 +2161,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2061,18 +2172,23 @@ def circuit(a): assert isinstance(jac[1], torch.Tensor) assert jac[1].shape == (4, 2) - def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, shots, device_vjp - ): + def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2080,7 +2196,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2088,7 +2203,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + hess = hessian(circuit, (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2108,7 +2223,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2116,6 +2231,13 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( @@ -2124,7 +2246,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2132,14 +2253,12 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a: circuit(_a, shots=shots), params) + hess = hessian(circuit, params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots - ): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): """The hessian of a single measurement with multiple params returns a tuple of arrays.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2148,8 +2267,15 @@ def test_hessian_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2157,7 +2283,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2165,7 +2290,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a, _b: circuit(_a, _b, shots=shots), (par_0, par_1)) + hess = hessian(circuit, (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2185,7 +2310,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2195,7 +2320,9 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - params = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) + dev = qml.device(dev_name, wires=2, shots=shots) + + params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( dev, @@ -2203,7 +2330,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2211,13 +2337,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a: circuit(_a, shots=shots), params) + hess = hessian(circuit, params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2228,8 +2354,15 @@ def test_hessian_probs_expval_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2237,7 +2370,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2246,7 +2378,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y, shots=shots)) + return torch.hstack(circuit(x, y)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2270,9 +2402,10 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2281,7 +2414,7 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) + par = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( dev, @@ -2289,7 +2422,6 @@ def test_hessian_expval_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2298,7 +2430,7 @@ def circuit(x): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x, shots=shots)) + return torch.hstack(circuit(x)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2308,9 +2440,10 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2328,7 +2461,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2337,7 +2469,7 @@ def circuit(x, y): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y, shots=shots)) + return torch.hstack(circuit(x, y)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2361,9 +2493,10 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_var_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2380,7 +2513,6 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2389,7 +2521,7 @@ def circuit(x): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x, shots=shots)) + return torch.hstack(circuit(x)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2399,11 +2531,14 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="torch") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="torch") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py index d8180ff7bd9..2a6ee306508 100644 --- a/tests/interfaces/test_autograd.py +++ b/tests/interfaces/test_autograd.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,310 +11,31 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the autograd interface""" -# pylint: disable=protected-access,too-few-public-methods -import sys - +"""Autograd specific tests for execute and default qubit 2.""" import autograd import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml +from pennylane import execute from pennylane import numpy as np -from pennylane.devices import DefaultQubitLegacy -from pennylane.gradients import finite_diff, param_shift -from pennylane.operation import AnyWires, Observable +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift +from pennylane.measurements import Shots pytestmark = pytest.mark.autograd -class TestAutogradExecuteUnitTests: - """Unit tests for autograd execution""" - - def test_import_error(self, mocker): - """Test that an exception is caught on import error""" - - mock = mocker.patch.object(autograd.extend, "defvjp") - mock.side_effect = ImportError() - - try: - del sys.modules["pennylane.workflow.interfaces.autograd"] - except KeyError: - pass - - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - - with qml.queuing.AnnotatedQueue() as q: - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.raises( - qml.QuantumFunctionError, - match="autograd not found. Please install the latest version " - "of autograd to enable the 'autograd' interface", - ): - qml.execute([tape], dev, gradient_fn=param_shift, interface="autograd") - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - device, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - )[0] - - qml.jacobian(cost)(a, device=dev) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if a gradient transform - is used with grad_on_execution=True""" - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, gradient_fn=param_shift, grad_on_execution=True)[0] - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - qml.jacobian(cost)(a, device=dev) - - def test_unknown_interface(self): - """Test that an error is raised if the interface is unknown""" - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, gradient_fn=param_shift, interface="None")[0] - - with pytest.raises(ValueError, match="interface must be in"): - cost(a, device=dev) - - def test_grad_on_execution(self, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - - a = np.array([0.1, 0.2], requires_grad=True) - cost(a) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_no_gradients_on_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - - a = np.array([0.1, 0.2], requires_grad=True) - cost(a) - - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - - qml.jacobian(cost)(a) - spy_gradients.assert_called() - - -class TestBatchTransformExecution: - """Tests to ensure batch transforms can be correctly executed - via qml.execute and batch_transform""" - - def test_no_batch_transform(self, mocker): - """Test that batch transforms can be disabled and enabled""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100000) - - H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0) - x = 0.6 - y = 0.2 - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=0) - qml.RY(y, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.spy(dev, "batch_transform") - - if not qml.operation.active_new_opmath(): - with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"): - _ = qml.execute([tape], dev, None, device_batch_transform=False) - else: - res = qml.execute([tape], dev, None, device_batch_transform=False) - assert np.allclose(res[0], np.cos(y), atol=0.1) - - spy.assert_not_called() - - res = qml.execute([tape], dev, None, device_batch_transform=True) - spy.assert_called() - - assert qml.math.shape(res[0]) == () - assert np.allclose(res[0], np.cos(y), rtol=0.05) - - def test_batch_transform_dynamic_shots(self): - """Tests that the batch transform considers the number of shots for the execution, not those - statically on the device.""" - dev = qml.device("default.qubit.legacy", wires=1) - H = 2.0 * qml.PauliZ(0) - qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)]) - res = qml.execute([qscript], dev, interface=None, override_shots=10) - assert res == (2.0,) - - +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cachesize=cachesize)[0] - - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cachesize=2) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] - - custom_cache = {} - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_caching_param_shift(self, tol): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum.""" - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] - - # Without caching, 5 jacobians should still be performed - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cache=None) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = qml.jacobian(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 - - # Check that calling the cost function again - # continues to evaluate the device (that is, the cache - # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 - assert np.allclose(grad1, grad2, atol=tol, rtol=0) - - # Check that calling the cost function again - # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 - assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + """Tests for caching behaviour""" @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): + def test_caching_param_shift_hessian(self, num_params): """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() params = np.arange(1, num_params + 1) / 10 N = len(params) @@ -328,13 +49,16 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache, max_diff=2)[0] + return qml.execute( + [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 + )[0] # No caching: number of executions is not ideal - hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) + with qml.Tracker(dev) as tracker: + hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) if num_params == 2: # compare to theoretical result @@ -345,7 +69,7 @@ def cost(x, cache): [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert np.allclose(expected, hess1) expected_runs = 1 # forward pass @@ -360,62 +84,26 @@ def cost(x, cache): # Each tape used to compute the Jacobian is then shifted again expected_runs += runs_for_jacobian * num_shifted_evals - assert dev.num_executions == expected_runs + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + with qml.Tracker(dev) as tracker2: + hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal + assert tracker2.totals["executions"] == expected_runs_ideal assert expected_runs_ideal < expected_runs - def test_caching_adjoint_no_grad_on_execution(self): - """Test that caching reduces the number of adjoint evaluations - when the grads is not on execution.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = np.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack( - qml.execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - ) - - # no caching, but jac for each batch still stored. - qml.jacobian(cost)(params, cache=None) - assert dev.num_executions == 2 - - # With caching, only 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = qml.jacobian(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 - def test_single_backward_pass_batch(self): """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift is requested for multiple tapes.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit") def f(x): tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) @@ -429,241 +117,293 @@ def f(x): out = qml.jacobian(f)(x) assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["batch_len"] == [2, 4] + assert dev.tracker.history["simulations"] == [1, 1, 1, 1, 1, 1] expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] assert qml.math.allclose(out, expected) - def test_single_backward_pass_split_hamiltonian(self): - """Tests that the backward pass is one single batch, not a bunch of batches, when parameter - shift derivatives are requested for a tape that the device split into batches.""" - - dev = qml.device("default.qubit.legacy", wires=2, shots=50000) - H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc") - - def f(x): - tape = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.expval(H)]) - return qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift)[0] +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), + ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, + Shots(None), + DefaultQubit(), + ), + ( + { + "gradient_fn": "adjoint", + "grad_on_execution": False, + "device_vjp": False, + }, + Shots(None), + DefaultQubit(), + ), + ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "device", "device_vjp": False}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(seed=904747894), + ), + ( + {"gradient_fn": "device", "device_vjp": True}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(seed=10490244), + ), +] - x = qml.numpy.array(0.1) - with dev.tracker: - out = qml.grad(f)(x) - assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["batch_len"] == [2, 4] - - assert qml.math.allclose(out, -np.cos(x) - np.sin(x), atol=0.05) - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestAutogradExecuteIntegration: """Test the autograd interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) def cost(a, b): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], dev, **execute_kwargs) + return execute([tape1, tape2], device, **execute_kwargs) a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) - res = cost(a, b) + with device.tracker: + res = cost(a, b) + + if execute_kwargs.get("grad_on_execution", False): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" a = np.array(0.1, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, **execute_kwargs)[0] + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - res = qml.jacobian(cost)(a) - assert res.shape == () + if shots.has_partitioned_shots: + res = qml.jacobian(lambda x: qml.math.hstack(cost(x)))(a) + else: + res = qml.jacobian(cost)(a) + assert res.shape == () # pylint: disable=no-member # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -np.sin(a), atol=atol_for_shots(shots)) - def test_jacobian(self, execute_kwargs, tol): + def test_jacobian(self, execute_kwargs, shots, device): """Test jacobian calculation""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - def cost(a, b, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) - - dev = qml.device("default.qubit.legacy", wires=2) + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - res = cost(a, b, device=dev) + res = cost(a, b) expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = qml.jacobian(cost)(a, b, device=dev) + res = qml.jacobian(cost)(a, b) assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (2,) - assert res[1].shape == (2,) + if shots.has_partitioned_shots: + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + for _r, _e in zip(res, expected): + assert np.allclose(_r[:2], _e, atol=atol_for_shots(shots)) + assert np.allclose(_r[2:], _e, atol=atol_for_shots(shots)) + else: + assert res[0].shape == (2,) + assert res[1].shape == (2,) - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected)) + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + for _r, _e in zip(res, expected): + assert np.allclose(_r, _e, atol=atol_for_shots(shots)) @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tape_no_parameters(self, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - dev = qml.device("default.qubit.legacy", wires=2) + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - def cost(params): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape3 = qml.tape.QuantumScript.from_queue(q3) - with qml.queuing.AnnotatedQueue() as q4: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.probs(wires=[0, 1]) - - tape4 = qml.tape.QuantumScript.from_queue(q4) - return sum( - autograd.numpy.hstack( - qml.execute([tape1, tape2, tape3, tape4], dev, **execute_kwargs) - ) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, ) + tape4 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, + ) + res = qml.execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = tuple(i for r in res for i in r) + return sum(autograd.numpy.hstack(res)) + params = np.array([0.1, 0.2], requires_grad=True) x, y = params res = cost(params) expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + expected = shots.num_copies * expected + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) grad = qml.grad(cost)(params) - expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + expected = np.array([-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]) + if shots.has_partitioned_shots: + expected = shots.num_copies * expected + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tapes_with_different_return_size(self, execute_kwargs): + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): """Test that tapes wit different can be executed and differentiated.""" - dev = qml.device("default.qubit.legacy", wires=2) + + if execute_kwargs["gradient_fn"] == "backprop": + pytest.xfail("backprop is not compatible with something about this situation.") def cost(params): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape3 = qml.tape.QuantumScript.from_queue(q3) - return autograd.numpy.hstack(qml.execute([tape1, tape2, tape3], dev, **execute_kwargs)) + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) + + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = execute([tape1, tape2, tape3], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = tuple(i for r in res for i in r) + return autograd.numpy.hstack(res) params = np.array([0.1, 0.2], requires_grad=True) + x, y = params res = cost(params) assert isinstance(res, np.ndarray) - assert res.shape == (4,) - + if not shots: + assert res.shape == (4,) + + if shots.has_partitioned_shots: + for i in (0, 1): + assert np.allclose(res[2 * i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[2 * i + 1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[4 + i], np.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[6 + i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], np.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + + if shots.has_partitioned_shots: + pytest.xfail("autograd jacobians do not work with ragged results and shot vectors.") + # TODO: autograd jacobians with ragged results and shot vectors jac = qml.jacobian(cost)(params) assert isinstance(jac, np.ndarray) - assert jac.shape == (4, 2) + if not shots.has_partitioned_shots: + assert jac.shape == (4, 2) # pylint: disable=no-member + + d1 = -np.sin(x) * np.cos(y) + d2 = -np.cos(x) * np.sin(y) + + if shots.has_partitioned_shots: + assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) + assert np.allclose(jac[3:4], 0, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) + + assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) + + else: + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - def test_reusing_quantum_tape(self, execute_kwargs, tol): + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): """Test re-using a quantum tape by passing new parameters""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + ) assert tape.trainable_params == [0, 1] def cost(a, b): new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return autograd.numpy.hstack(qml.execute([new_tape], dev, **execute_kwargs)[0]) + return autograd.numpy.hstack(execute([new_tape], device, **execute_kwargs)[0]) jac_fn = qml.jacobian(cost) jac = jac_fn(a, b) @@ -675,7 +415,7 @@ def cost(a, b): # values of the parameters for subsequent calls res2 = cost(2 * a, b) expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) jac = jac_fn(a, b) @@ -684,92 +424,89 @@ def cost(a, b): [0, -np.cos(2 * a) * np.cos(b)], ) assert isinstance(jac, tuple) and len(jac) == 2 - assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected)) + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing(self, execute_kwargs): + def test_classical_processing(self, execute_kwargs, device, shots): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) - def cost(a, b, c, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + np.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + np.sin(a), wires=0), + ] - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + if shots.has_partitioned_shots: + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0]) + return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = qml.jacobian(cost)(a, b, c, device=dev) + res = qml.jacobian(cost)(a, b, c) # Only two arguments are trainable assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () + + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - def test_no_trainable_parameters(self, execute_kwargs): + def test_no_trainable_parameters(self, execute_kwargs, shots, device): """Test evaluation and Jacobian if there are no trainable parameters""" a = np.array(0.1, requires_grad=False) b = np.array(0.2, requires_grad=False) - def cost(a, b, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, b, device=dev) - assert res.shape == (2,) + res = cost(a, b) + assert res.shape == (2 * shots.num_copies,) if shots else (2,) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.jacobian(cost)(a, b, device=dev) + res = qml.jacobian(cost)(a, b) assert len(res) == 0 def loss(a, b): - return np.sum(cost(a, b, device=dev)) + return np.sum(cost(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): res = qml.grad(loss)(a, b) assert np.allclose(res, 0) - def test_matrix_parameter(self, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) - def cost(a, U, device): - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, U, device=dev) - assert isinstance(res, np.ndarray) - assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + res = cost(a, U) + assert np.allclose(res, -np.cos(a), atol=atol_for_shots(shots)) jac_fn = qml.jacobian(cost) - jac = jac_fn(a, U, device=dev) + jac = jac_fn(a, U) assert isinstance(jac, np.ndarray) - assert np.allclose(jac, np.sin(a), atol=tol, rtol=0) + assert np.allclose(jac, np.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): + def test_differentiable_expand(self, execute_kwargs, device, shots): """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -778,28 +515,37 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - def cost_fn(a, p, device): - with qml.queuing.AnnotatedQueue() as q_tape: - qml.RX(a, wires=0) - U3(*p, wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q_tape) - tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) - return qml.execute([tape], device, **execute_kwargs)[0] + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=1) - res = cost_fn(a, p, device=dev) + res = cost_fn(a, p) expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost_fn) - res = jac_fn(a, p, device=dev) + res = jac_fn(a, p) expected = np.array( [ np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), @@ -810,31 +556,22 @@ def cost_fn(a, p, device): * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - - def cost(x, y, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - dev = qml.device("default.qubit.legacy", wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) - res = cost(x, y, device=dev) + res = cost(x, y) expected = np.array( [ [ @@ -845,10 +582,10 @@ def cost(x, y, device): ], ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost) - res = jac_fn(x, y, device=dev) + res = jac_fn(x, y) assert isinstance(res, tuple) and len(res) == 2 assert res[0].shape == (4,) assert res[1].shape == (4,) @@ -871,38 +608,30 @@ def cost(x, y, device): ), ) - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - def test_ragged_differentiation(self, execute_kwargs, tol): + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - def cost(x, y, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) - - dev = qml.device("default.qubit.legacy", wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) - res = cost(x, y, device=dev) + res = cost(x, y) expected = np.array( [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost) - res = jac_fn(x, y, device=dev) + res = jac_fn(x, y) assert isinstance(res, tuple) and len(res) == 2 assert res[0].shape == (3,) assert res[1].shape == (3,) @@ -911,34 +640,8 @@ def cost(x, y, device): np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), ) - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - def test_sampling(self, execute_kwargs): - """Test sampling works as expected""" - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - - def cost(device): - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] - - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - res = cost(device=dev) - assert isinstance(res, tuple) - assert len(res) == 2 - assert res[0].shape == (shots,) - assert res[1].shape == (shots,) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) class TestHigherOrderDerivatives: @@ -955,24 +658,15 @@ class TestHigherOrderDerivatives: def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] res = cost_fn(params) @@ -995,83 +689,20 @@ def cost_fn(x): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_adjoint_hessian_one_param(self, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.autograd", wires=2) - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.grad(cost_fn))(params) - - assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0) - - def test_adjoint_hessian_multiple_params(self, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.autograd", wires=2) - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack( - qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - ) - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.jacobian(cost_fn))(params) - - assert np.allclose(res, np.zeros([2, 2, 2]), atol=tol, rtol=0) - def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) return result[0] + result[1][0] res = cost_fn(params) @@ -1092,149 +723,26 @@ def cost_fn(x): assert np.allclose(res, expected, atol=tol, rtol=0) -class TestOverridingShots: - """Test overriding shots on execution""" - - def test_changing_shots(self, mocker, tol): - """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.spy(dev, "sample") - - # execute with device default shots (None) - res = qml.execute([tape], dev, gradient_fn=param_shift) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() - - # execute with shots=100 - res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = qml.execute([tape], dev, gradient_fn=param_shift) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # same single call from above, no additional calls - - def test_overriding_shots_with_same_value(self, mocker): - """Overriding shots with the same value as the device will have no effect""" - dev = qml.device("default.qubit.legacy", wires=2, shots=123) - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.Mock(wraps=qml.Device.shots.fset) - # pylint:disable=assignment-from-no-return,too-many-function-args - mock_property = qml.Device.shots.setter(spy) - mocker.patch.object(qml.Device, "shots", mock_property) - - qml.execute([tape], dev, gradient_fn=param_shift, override_shots=123) - # overriden shots is the same, no change - spy.assert_not_called() - - qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) - # overriden shots is not the same, shots were changed - spy.assert_called() - - # shots were temporarily set to the overriden value - assert spy.call_args_list[0][0] == (dev, 100) - # shots were then returned to the built-in value - assert spy.call_args_list[1][0] == (dev, 123) - - def test_overriding_device_with_shot_vector(self): - """Overriding a device that has a batch of shots set - results in original shots being returned after execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=[10, (1, 3), 5]) - - assert dev.shots == 18 - assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] - - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100)[0] - - assert isinstance(res, np.ndarray) - assert res.shape == () - - # device is unchanged - assert dev.shots == 18 - assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] - - res = qml.execute([tape], dev, gradient_fn=param_shift)[0] - assert len(res) == 5 - - @pytest.mark.xfail(reason="Shots vector must be adapted for new return types.") - def test_gradient_integration(self): - """Test that temporarily setting the shots works - for gradient computations""" - # TODO: Update here when shot vectors are supported - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - a, b = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(a, b, shots): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, override_shots=shots)[0] - - res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000]) - assert dev.shots is None - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] - assert all( - np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected) - ) - - -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift}, - {"gradient_fn": finite_diff}, -] - - -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -1243,8 +751,8 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) return _cost_fn @@ -1276,246 +784,56 @@ def cost_fn_jacobian(weights, coeffs1, coeffs2): ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): """Test hamiltonian with no trainable parameters.""" - # pylint: disable=unused-argument + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) coeffs2 = np.array([0.7], requires_grad=False) weights = np.array([0.4, 0.5], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev) + res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): """Test hamiltonian with trainable parameters.""" - # pylint: disable=unused-argument + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) coeffs2 = np.array([0.7], requires_grad=True) weights = np.array([0.4, 0.5], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)) + res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2)) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -class TestCustomJacobian: - """Test for custom Jacobian.""" - - def test_custom_jacobians(self): - """Test custom Jacobian device methood""" - - class CustomJacobianDevice(DefaultQubitLegacy): - @classmethod - def capabilities(cls): - capabilities = super().capabilities() - capabilities["provides_jacobian"] = True - return capabilities - - def jacobian(self, tape): - # pylint: disable=unused-argument - return np.array([1.0, 2.0, 3.0, 4.0]) - - dev = CustomJacobianDevice(wires=2) - - @qml.qnode(dev, diff_method="device") - def circuit(v): - qml.RX(v, wires=0) - return qml.probs(wires=[0, 1]) - - d_circuit = qml.jacobian(circuit, argnum=0) - - params = np.array(1.0, requires_grad=True) - - d_out = d_circuit(params) - assert np.allclose(d_out, np.array([1.0, 2.0, 3.0, 4.0])) - - def test_custom_jacobians_param_shift(self): - """Test computing the gradient using the parameter-shift - rule with a device that provides a jacobian""" - - class MyQubit(DefaultQubitLegacy): - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update( - provides_jacobian=True, - ) - return capabilities - - def jacobian(self, *args, **kwargs): - raise NotImplementedError() - - dev = MyQubit(wires=2) - - @qml.qnode(dev, diff_method="parameter-shift", grad_on_execution=False) - def qnode(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - def cost(a, b): - return autograd.numpy.hstack(qnode(a, b)) - - res = qml.jacobian(cost)(a, b) - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - - assert np.allclose(res[0], expected[0]) - assert np.allclose(res[1], expected[1]) - - -class SpecialObject: - """SpecialObject - A special object that conveniently encapsulates the return value of - a special observable supported by a special device and which supports - multiplication with scalars and addition. - """ - - def __init__(self, val): - self.val = val - - def __mul__(self, other): - new = SpecialObject(self.val) - new *= other - return new - - def __imul__(self, other): - self.val *= other - return self - - def __rmul__(self, other): - return self * other - - def __iadd__(self, other): - self.val += other.val if isinstance(other, self.__class__) else other - return self - - def __add__(self, other): - new = SpecialObject(self.val) - new += other.val if isinstance(other, self.__class__) else other - return new - - def __radd__(self, other): - return self + other - - -class SpecialObservable(Observable): - """SpecialObservable""" - - num_wires = AnyWires - num_params = 0 - par_domain = None - - def diagonalizing_gates(self): - """Diagonalizing gates""" - return [] - - -class DeviceSupportingSpecialObservable(DefaultQubitLegacy): - name = "Device supporting SpecialObservable" - short_name = "default.qubit.specialobservable" - observables = DefaultQubitLegacy.observables.union({"SpecialObservable"}) - - @staticmethod - def _asarray(arr, dtype=None): - # pylint: disable=unused-argument - return arr - - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update( - provides_jacobian=True, - ) - return capabilities - - def expval(self, observable, **kwargs): - if self.analytic and isinstance(observable, SpecialObservable): - val = super().expval(qml.PauliZ(wires=0), **kwargs) - return np.array(SpecialObject(val)) - - return super().expval(observable, **kwargs) - - def jacobian(self, tape): - # we actually let pennylane do the work of computing the - # jacobian for us but return it as a device jacobian - gradient_tapes, fn = qml.gradients.param_shift(tape) - tape_jacobian = fn(qml.execute(gradient_tapes, self, None)) - return tape_jacobian - - -@pytest.mark.autograd -class TestObservableWithObjectReturnType: - """Unit tests for qnode returning a custom object""" - - def test_custom_return_type(self): - """Test custom return values for a qnode""" - - dev = DeviceSupportingSpecialObservable(wires=1, shots=None) - - # force diff_method='parameter-shift' because otherwise - # PennyLane swaps out dev for default.qubit.autograd - @qml.qnode(dev, diff_method="parameter-shift") - def qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - @qml.qnode(dev, diff_method="parameter-shift") - def reference_qnode(x): - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(wires=0)) - - out = qnode(0.2) - assert isinstance(out, np.ndarray) - assert isinstance(out.item(), SpecialObject) - assert np.isclose(out.item().val, reference_qnode(0.2)) - - def test_jacobian_with_custom_return_type(self): - """Test differentiation of a QNode on a device supporting a - special observable that returns an object rather than a number.""" - - dev = DeviceSupportingSpecialObservable(wires=1, shots=None) - - # force diff_method='parameter-shift' because otherwise - # PennyLane swaps out dev for default.qubit.autograd - @qml.qnode(dev, diff_method="parameter-shift") - def qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - @qml.qnode(dev, diff_method="parameter-shift") - def reference_qnode(x): - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(wires=0)) - - reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),) - - assert np.isclose( - reference_jac, - qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val, - ) - - # now check that also the device jacobian works with a custom return type - @qml.qnode(dev, diff_method="device") - def device_gradient_qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - assert np.isclose( - reference_jac, - qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val, - ) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_autograd_qnode.py b/tests/interfaces/test_autograd_qnode.py index 818d7d58ff5..a0d1e0f2e66 100644 --- a/tests/interfaces/test_autograd_qnode.py +++ b/tests/interfaces/test_autograd_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2020 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the autograd interface with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison, -# pylint: disable=unnecessary-lambda-assignment +# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-dict-literal, no-name-in-module + import autograd import autograd.numpy as anp import pytest @@ -21,32 +21,47 @@ import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit +from tests.param_shift_dev import ParamShiftDerivativesDevice +# dev, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [qml.device("default.qubit"), "finite-diff", False, False], + [qml.device("default.qubit"), "parameter-shift", False, False], + [qml.device("default.qubit"), "backprop", True, False], + [qml.device("default.qubit"), "adjoint", True, False], + [qml.device("default.qubit"), "adjoint", False, False], + [qml.device("default.qubit"), "spsa", False, False], + [qml.device("default.qubit"), "hadamard", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, False], + [ParamShiftDerivativesDevice(), "best", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, True], ] interface_qubit_device_and_diff_method = [ - ["autograd", "default.qubit.legacy", "finite-diff", False], - ["autograd", "default.qubit.legacy", "parameter-shift", False], - ["autograd", "default.qubit.legacy", "backprop", True], - ["autograd", "default.qubit.legacy", "adjoint", True], - ["autograd", "default.qubit.legacy", "adjoint", False], - ["autograd", "default.qubit.legacy", "spsa", False], - ["autograd", "default.qubit.legacy", "hadamard", False], - ["auto", "default.qubit.legacy", "finite-diff", False], - ["auto", "default.qubit.legacy", "parameter-shift", False], - ["auto", "default.qubit.legacy", "backprop", True], - ["auto", "default.qubit.legacy", "adjoint", True], - ["auto", "default.qubit.legacy", "adjoint", False], - ["auto", "default.qubit.legacy", "spsa", False], - ["auto", "default.qubit.legacy", "hadamard", False], + ["autograd", DefaultQubit(), "finite-diff", False, False], + ["autograd", DefaultQubit(), "parameter-shift", False, False], + ["autograd", DefaultQubit(), "backprop", True, False], + ["autograd", DefaultQubit(), "adjoint", True, False], + ["autograd", DefaultQubit(), "adjoint", False, False], + ["autograd", DefaultQubit(), "adjoint", True, True], + ["autograd", DefaultQubit(), "adjoint", False, True], + ["autograd", DefaultQubit(), "spsa", False, False], + ["autograd", DefaultQubit(), "hadamard", False, False], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, True], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, True], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, False], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, False], + ["auto", DefaultQubit(), "finite-diff", False, False], + ["auto", DefaultQubit(), "parameter-shift", False, False], + ["auto", DefaultQubit(), "backprop", True, False], + ["auto", DefaultQubit(), "adjoint", True, False], + ["auto", DefaultQubit(), "adjoint", False, False], + ["auto", DefaultQubit(), "spsa", False, False], + ["auto", DefaultQubit(), "hadamard", False, False], + ["auto", qml.device("lightning.qubit", wires=5), "adjoint", False, False], + ["auto", qml.device("lightning.qubit", wires=5), "adjoint", True, False], ] pytestmark = pytest.mark.autograd @@ -57,78 +72,21 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with Autograd integrates with the PennyLane stack""" - # pylint: disable=unused-argument - - def test_nondiff_param_unwrapping( - self, interface, dev_name, diff_method, grad_on_execution, mocker + # pylint:disable=unused-argument + def test_execution_no_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp ): - """Test that non-differentiable parameters are correctly unwrapped - to NumPy ndarrays or floats (if 0-dimensional)""" - if diff_method != "parameter-shift": - pytest.skip("Test only supports parameter-shift") - - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method=diff_method) - def circuit(x, y): - qml.RX(x[0], wires=0) - qml.Rot(*x[1:], wires=0) - qml.RY(y[0], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False) - y = np.array([0.5], requires_grad=True) - - param_data = [] - - def mock_apply(*args, **kwargs): - for op in args[0]: - param_data.extend(op.data) - - mocker.patch.object(dev, "apply", side_effect=mock_apply) - circuit(x, y) - assert param_data == [0.1, 0.2, 0.3, 0.4, 0.5] - assert not any(isinstance(p, np.tensor) for p in param_data) - - # test the jacobian works correctly - param_data = [] - qml.grad(circuit)(x, y) - assert param_data == [ - 0.1, - 0.2, - 0.3, - 0.4, - 0.5, - 0.1, - 0.2, - 0.3, - 0.4, - 0.5 + np.pi / 2, - 0.1, - 0.2, - 0.3, - 0.4, - 0.5 - np.pi / 2, - ] - assert not any(isinstance(p, np.tensor) for p in param_data) - - def test_execution_no_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works without an interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface=None, diff_method=diff_method) + @qnode(dev, interface=None) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -138,26 +96,24 @@ def circuit(a): res = circuit(a) - # without the interface, the QNode simply returns a scalar array - assert isinstance(res, np.ndarray) - assert res.shape == tuple() - - # gradients should cause an error - with pytest.raises(TypeError, match="must be real number, not ArrayBox"): - qml.grad(circuit)(a) + # without the interface, the QNode simply returns a scalar array or float + assert isinstance(res, (np.ndarray, float)) + assert qml.math.shape(res) == tuple() # pylint: disable=comparison-with-callable - def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): + def test_execution_with_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -165,31 +121,27 @@ def circuit(a): a = np.array(0.1, requires_grad=True) assert circuit.interface == interface - # gradients should work grad = qml.grad(circuit)(a) - assert isinstance(grad, float) assert grad.shape == tuple() - def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol, device_vjp): """Test jacobian calculation""" - num_wires = 2 kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) + if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -220,24 +172,23 @@ def cost(x, y): assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian_no_evaluate( + self, interface, dev, diff_method, grad_on_execution, tol, device_vjp + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - num_wires = 2 kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -263,36 +214,31 @@ def cost(x, y): assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test setting jacobian options""" - wires = [0] - if diff_method in ["backprop", "adjoint"]: - pytest.skip("Test does not support backprop or adjoint method") - elif diff_method == "finite-diff": - kwargs = {"h": 1e-8, "approx_order": 2} - elif diff_method == "parameter-shift": - kwargs = {"shifts": [(0.1,), (0.2,)]} - elif diff_method == "hadamard": - wires = [0, "aux"] - kwargs = {"aux_wire": qml.wires.Wires("aux"), "device_wires": wires} - else: - kwargs = {} + if diff_method != "finite-diff": + pytest.skip("Test only supports finite diff.") a = np.array([0.1, 0.2], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=wires) - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + @qnode( + dev, + interface=interface, + h=1e-8, + order=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - circuit(a) - qml.jacobian(circuit)(a) - def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_changing_trainability( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -301,9 +247,13 @@ def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_e a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method="parameter-shift", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -340,21 +290,18 @@ def loss(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): + def test_classical_processing(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -372,12 +319,17 @@ def circuit(a, b, c): assert res[0].shape == () assert res[1].shape == () - def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): + def test_no_trainable_parameters( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test evaluation and Jacobian if there are no trainable parameters""" - dev = qml.device(dev_name, wires=2) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -396,35 +348,34 @@ def circuit(a, b): assert len(res) == 2 assert isinstance(res, tuple) - def cost0(x, y): + def cost(x, y): return autograd.numpy.hstack(circuit(x, y)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - assert not qml.jacobian(cost0)(a, b) + assert not qml.jacobian(cost)(a, b) - def cost1(a, b): + def cost2(a, b): return np.sum(circuit(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - grad = qml.grad(cost1)(a, b) + grad = qml.grad(cost2)(a, b) assert grad == tuple() - def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_matrix_parameter( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -440,13 +391,18 @@ def circuit(U, a): assert np.allclose(res, np.sin(a), atol=tol, rtol=0) def test_gradient_non_differentiable_exception( - self, interface, dev_name, diff_method, grad_on_execution + self, interface, dev, diff_method, grad_on_execution, device_vjp ): """Test that an exception is raised if non-differentiable data is differentiated""" - dev = qml.device(dev_name, wires=2) - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(data1): qml.templates.AmplitudeEmbedding(data1, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @@ -457,18 +413,27 @@ def circuit(data1): with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"): grad_fn(data1) - def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_differentiable_expand( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that operation and nested tape expansion is differentiable""" kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) + if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=20) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA + # pylint: disable=too-few-public-methods class U3(qml.U3): + """Custom U3.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -477,7 +442,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device(dev_name, wires=2) a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) @@ -491,14 +455,14 @@ def circuit(a, p): expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) res = qml.grad(circuit)(a, p) - assert isinstance(res, np.ndarray) - assert len(res) == 3 + # assert isinstance(res, np.ndarray) + # assert len(res) == 3 expected = np.array( [ @@ -517,9 +481,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and remains differentiable.""" - def test_changing_shots(self, mocker, tol): + def test_changing_shots(self): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -527,31 +491,21 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + res = circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # same single call performed above + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # pylint: disable=comparison-with-callable @pytest.mark.xfail(reason="Param shift and shot vectors.") def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -574,62 +528,57 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = np.array([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(dev) + @qnode(DefaultQubit()) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b) + assert cost_fn.gradient_fn == "backprop" # gets restored to default + + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - # pylint: disable=unused-argument - def test_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with a single prob output""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") + kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -650,25 +599,23 @@ def circuit(x, y): assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) def test_multiple_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -717,22 +664,24 @@ def cost(x, y): ) assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_ragged_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") + kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - num_wires = 2 - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -768,22 +717,25 @@ def cost(x, y): assert np.allclose(res[1], expected[1], atol=tol, rtol=0) def test_ragged_differentiation_variance( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -804,12 +756,12 @@ def circuit(x, y): assert isinstance(res, tuple) assert len(res) == 2 - assert isinstance(res[0], np.ndarray) - assert res[0].shape == () + # assert isinstance(res[0], np.ndarray) + # assert res[0].shape == () assert np.allclose(res[0], expected_var, atol=tol, rtol=0) - assert isinstance(res[1], np.ndarray) - assert res[1].shape == (2,) + # assert isinstance(res[1], np.ndarray) + # assert res[1].shape == (2,) assert np.allclose(res[1], expected_probs, atol=tol, rtol=0) def cost(x, y): @@ -826,33 +778,38 @@ def cost(x, y): assert isinstance(jac, tuple) assert len(jac) == 2 - assert isinstance(jac[0], np.ndarray) - assert jac[0].shape == (3,) + # assert isinstance(jac[0], np.ndarray) + # assert jac[0].shape == (3,) assert np.allclose(jac[0], expected[0], atol=tol, rtol=0) - assert isinstance(jac[1], np.ndarray) - assert jac[1].shape == (3,) + # assert isinstance(jac[1], np.ndarray) + # assert jac[1].shape == (3,) assert np.allclose(jac[1], expected[1], atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + # pylint: disable=too-few-public-methods class Template(qml.templates.StronglyEntanglingLayers): + """Custom template.""" + def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit2(data, weights): qml.templates.AngleEmbedding(data, wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -876,19 +833,22 @@ def cost(w1, w2): assert len(res) == 2 - def test_chained_gradient_value(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_chained_gradient_value( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the returned gradient value for two chained qubit QNodes is correct.""" - kwargs = dict(interface=interface, diff_method=diff_method) + kwargs = dict( + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev1 = qml.device(dev_name, wires=num_wires) + dev1 = qml.device("default.qubit") @qnode(dev1, **kwargs) def circuit1(a, b, c): @@ -899,9 +859,11 @@ def circuit1(a, b, c): qml.CNOT(wires=[1, 2]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2)) - dev2 = qml.device("default.qubit.legacy", wires=num_wires) + dev2 = dev - @qnode(dev2, interface=interface, diff_method=diff_method) + @qnode( + dev2, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + ) def circuit2(data, weights): qml.RX(data[0], wires=0) qml.RX(data[1], wires=1) @@ -971,19 +933,20 @@ def cost(a, b, c, weights): # to the first parameter of circuit1. assert circuit1.qtape.trainable_params == [1, 2] - def test_second_derivative(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_second_derivative( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -1013,18 +976,17 @@ def circuit(x): assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1039,8 +1001,8 @@ def circuit(x): expected_res = np.cos(a) * np.cos(b) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected_res, atol=tol, rtol=0) grad_fn = qml.grad(circuit) @@ -1068,19 +1030,18 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_unused_parameter( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1109,19 +1070,20 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_vector_valued( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1137,7 +1099,7 @@ def circuit(x): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) + assert res.shape == (2,) # pylint: disable=comparison-with-callable assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1174,19 +1136,18 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1229,19 +1190,18 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(a, b): @@ -1255,7 +1215,7 @@ def circuit(a, b): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) + assert res.shape == (2,) # pylint: disable=comparison-with-callable assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1272,8 +1232,12 @@ def circuit(a, b): assert g[1].shape == (2,) assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0) - jac_fn_a = lambda *args: jac_fn(*args)[0] - jac_fn_b = lambda *args: jac_fn(*args)[1] + def jac_fn_a(*args): + return jac_fn(*args)[0] + + def jac_fn_b(*args): + return jac_fn(*args)[1] + hess_a = qml.jacobian(jac_fn_a)(a, b) hess_b = qml.jacobian(jac_fn_b)(a, b) assert isinstance(hess_a, tuple) and len(hess_a) == 2 @@ -1295,18 +1259,17 @@ def circuit(a, b): assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0) assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=2) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1320,14 +1283,18 @@ def circuit(x): a, b = x - cos_prod = np.cos(a) * np.cos(b) - expected_res = (cos_prod, [0.5 + 0.5 * cos_prod, 0.5 - 0.5 * cos_prod]) - res = circuit(x) - assert all(qml.math.allclose(r, e) for r, e in zip(res, expected_res)) + expected_res = [ + np.cos(a) * np.cos(b), + 0.5 + 0.5 * np.cos(a) * np.cos(b), + 0.5 - 0.5 * np.cos(a) * np.cos(b), + ] def cost_fn(x): return autograd.numpy.hstack(circuit(x)) + res = cost_fn(x) + assert qml.math.allclose(res, expected_res) + jac_fn = qml.jacobian(cost_fn) hess = qml.jacobian(jac_fn)(x) @@ -1351,18 +1318,21 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - dev = qml.device(dev_name, wires=2) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning does not support state adjoint diff.") x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1393,20 +1363,25 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): + def test_projector( + self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the variance of a projector is correctly returned""" + if diff_method == "adjoint": + pytest.skip("adjoint supports either expvals or diagonal measurements.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) P = np.array(state, requires_grad=False) x, y = np.array([0.765, -0.654], requires_grad=True) @@ -1419,8 +1394,8 @@ def circuit(x, y): res = circuit(x, y) expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) jac = qml.jacobian(circuit)(x, y) @@ -1444,106 +1419,60 @@ def circuit(x, y): assert np.allclose(jac, expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - r = np.array(0.543, requires_grad=True) - phi = np.array(-0.654, requires_grad=True) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - # circuit jacobians - res = qml.jacobian(circuit)(r, phi) - expected = np.array( - [ - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ] + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - tol = TOL_FOR_SPSA - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - n = np.array(0.12, requires_grad=True) - a = np.array(0.765, requires_grad=True) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = qml.jacobian(circuit)(n, a) - expected = np.array([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -def test_adjoint_reuse_device_state(mocker): - """Tests that the autograd interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - spy = mocker.spy(dev, "adjoint_jacobian") + phi = np.array(1.23, requires_grad=True) + theta = np.array(4.56, requires_grad=True) - qml.grad(circ, argnum=0)(1.0) - assert circ.device.num_executions == 1 + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - spy.assert_called_with(mocker.ANY, use_device_state=True) + gradient = qml.grad(circuit)(phi, theta) + exp_theta_grad = qml.grad(expected_circuit)(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method +) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Autograd interface""" @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff + self, dev, diff_method, grad_on_execution, max_diff, device_vjp ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" @@ -1552,20 +1481,22 @@ def test_gradient_expansion_trainable_only( if max_diff == 2 and diff_method == "hadamard": pytest.skip("Max diff > 1 not supported for Hadamard gradient.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - + # pylint: disable=too-few-public-methods class PhaseShift(qml.PhaseShift): + """dummy phase shift.""" + grad_method = None def decomposition(self): return [qml.RY(3 * self.data[0], wires=self.wires)] - @qnode(dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff) + @qnode( + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + max_diff=max_diff, + device_vjp=device_vjp, + ) def circuit(x, y): qml.Hadamard(wires=0) PhaseShift(x, wires=0) @@ -1576,11 +1507,11 @@ def circuit(x, y): y = np.array(0.7, requires_grad=False) circuit(x, y) - qml.grad(circuit)(x, y) + _ = qml.grad(circuit)(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol + self, dev, diff_method, grad_on_execution, max_diff, tol, device_vjp ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" @@ -1588,15 +1519,16 @@ def test_hamiltonian_expansion_analytic( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, ) + if diff_method in ["adjoint", "hadamard"]: pytest.skip("The diff method requested does not yet support Hamiltonians") elif diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1626,7 +1558,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method in ("parameter-shift", "backprop") and max_diff == 2: + if ( + diff_method in ("parameter-shift", "backprop") + and max_diff == 2 + and dev.name != "param_shift.qubit" + ): if diff_method == "backprop": with pytest.warns(UserWarning, match=r"Output seems independent of input."): grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) @@ -1644,28 +1580,26 @@ def circuit(data, weights, coeffs): @pytest.mark.slow @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, max_diff, device_vjp ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": - gradient_kwargs = dict( - h=H_FOR_SPSA, - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), - num_directions=20, - ) + gradient_kwargs = { + "h": H_FOR_SPSA, + "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), + "num_directions": 10, + } tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1673,6 +1607,7 @@ def test_hamiltonian_expansion_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1688,13 +1623,15 @@ def circuit(data, weights, coeffs): c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients - grad = qml.grad(circuit)(d, w, c) + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + + grad = qml.grad(circuit)(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1704,12 +1641,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - with pytest.warns(UserWarning, match=r"Output seems independent of input."): - grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) + if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": + grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c, shots=50000) assert np.allclose(grad2_c, 0, atol=tol) - grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c) + grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c, shots=50000) expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ 0, -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), @@ -1723,36 +1659,38 @@ class TestSample: def test_backprop_error(self): """Test that sampling in backpropagation grad_on_execution raises an error""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() @qnode(dev, diff_method="backprop") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - with pytest.raises(qml.QuantumFunctionError, match="only supported when shots=None"): - circuit(shots=10) # pylint: disable=unexpected-keyword-arg + with pytest.raises( + qml.QuantumFunctionError, match="does not support backprop with requested" + ): + circuit(shots=10) def test_sample_dimension(self): """Test that the sample function outputs samples of the right size""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=n_sample) assert isinstance(res, tuple) assert len(res) == 2 - assert res[0].shape == (10,) + assert res[0].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(res[0], np.ndarray) - assert res[1].shape == (10,) + assert res[1].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(res[1], np.ndarray) def test_sample_combination(self): @@ -1760,7 +1698,7 @@ def test_sample_combination(self): n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + dev = DefaultQubit() @qnode(dev, diff_method="parameter-shift") def circuit(): @@ -1768,29 +1706,29 @@ def circuit(): return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=n_sample) assert isinstance(result, tuple) assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert isinstance(result[1], np.ndarray) - assert isinstance(result[2], np.ndarray) + assert isinstance(result[1], (float, np.ndarray)) + assert isinstance(result[2], (float, np.ndarray)) assert result[0].dtype == np.dtype("float") def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=n_sample) assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) @@ -1800,44 +1738,44 @@ def test_multi_wire_sample_regular_shape(self): where a rectangular array is expected""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=n_sample) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) assert len(result) == 3 - assert result[0].shape == (10,) + assert result[0].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[0], np.ndarray) - assert result[1].shape == (10,) + assert result[1].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[1], np.ndarray) - assert result[2].shape == (10,) + assert result[2].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[2], np.ndarray) -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution,device_vjp", qubit_device_and_diff_method +) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - # pylint: disable=unused-argument - - def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution): + def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, device_vjp): """For one measurement and one param, the gradient is a float.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1847,25 +1785,20 @@ def circuit(a): grad = qml.grad(circuit)(a) - import sys + assert isinstance(grad, np.tensor if diff_method == "backprop" else float) - python_version = sys.version_info.minor - if diff_method == "backprop" and python_version > 7: - # Since numpy 1.23.0 - assert isinstance(grad, np.ndarray) - else: - assert isinstance(grad, float) - - def test_grad_single_measurement_multiple_param(self, dev_name, diff_method, grad_on_execution): + def test_grad_single_measurement_multiple_param( + self, dev, diff_method, grad_on_execution, device_vjp + ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1882,17 +1815,17 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1907,21 +1840,18 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1935,21 +1865,18 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1969,20 +1896,17 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single array is returned.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1995,20 +1919,17 @@ def circuit(a): assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a single params return an array.""" - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -2025,21 +1946,17 @@ def cost(x): assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2063,21 +1980,17 @@ def cost(x, y): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2093,14 +2006,8 @@ def cost(x): assert isinstance(jac, np.ndarray) assert jac.shape == (5, 2) - def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2108,7 +2015,14 @@ def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_exe par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2129,22 +2043,24 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_expval_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_expval_multiple_param_array( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of single measurement with a multiple params array return a single array.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2156,9 +2072,8 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2168,7 +2083,14 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2189,18 +2111,25 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_var_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_param_array( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2212,19 +2141,24 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_probs_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_probs_expval_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2249,18 +2183,23 @@ def cost(x, y): assert hess[1].shape == (6,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2273,11 +2212,12 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) + assert hess.shape == (3, 2, 2) # pylint: disable=no-member - def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_probs_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2287,7 +2227,14 @@ def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_ par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2311,18 +2258,25 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (6,) - def test_hessian_var_probs_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_param_array2( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2335,15 +2289,14 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) + assert hess.shape == (3, 2, 2) # pylint: disable=no-member -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) + dev = DefaultQubit() @qml.qnode(dev, interface="autograd") def circuit(): diff --git a/tests/interfaces/test_autograd_qnode_shot_vector.py b/tests/interfaces/test_autograd_qnode_shot_vector.py index 6aee767d01e..87e534fdb52 100644 --- a/tests/interfaces/test_autograd_qnode_shot_vector.py +++ b/tests/interfaces/test_autograd_qnode_shot_vector.py @@ -29,9 +29,9 @@ spsa_kwargs = {"h": 0.05, "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 0.05}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", spsa_kwargs], + ["default.qubit", "finite-diff", {"h": 0.05}], + ["default.qubit", "parameter-shift", {}], + ["default.qubit", "spsa", spsa_kwargs], ] TOLS = { diff --git a/tests/interfaces/test_execute.py b/tests/interfaces/test_execute.py index e12c5d8ff55..87201a27618 100644 --- a/tests/interfaces/test_execute.py +++ b/tests/interfaces/test_execute.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,18 +11,219 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -""" -Interface independent tests for qml.execute -""" - +"""Tests for exeuction with default qubit 2 independent of any interface.""" import pytest import pennylane as qml +from pennylane import numpy as np +from pennylane.devices import DefaultQubit +from pennylane.workflow.execution import _preprocess_expand_fn + + +class TestPreprocessExpandFn: + """Tests the _preprocess_expand_fn helper function.""" + + def test_provided_is_callable(self): + """Test that if the expand_fn is not "device", it is simply returned.""" + + dev = DefaultQubit() + + def f(tape): + return tape + + out = _preprocess_expand_fn(f, dev, 10) + assert out is f + + def test_new_device_blank_expand_fn(self): + """Test that the expand_fn is blank if is new device.""" + + dev = DefaultQubit() + + out = _preprocess_expand_fn("device", dev, 10) + + x = [1] + assert out(x) is x + + +class TestBatchTransformHelper: + """Unit tests for the _batch_transform helper function.""" + + def test_warns_if_requested_off(self): + """Test that a warning is raised if the the batch transform is requested to not be used.""" + + # pylint: disable=too-few-public-methods + class CustomOp(qml.operation.Operator): + """Dummy operator.""" + + def decomposition(self): + return [qml.PauliX(self.wires[0])] + + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): + program, _ = dev.preprocess() + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + qml.execute( + (qs, qs), device=dev, device_batch_transform=False, transform_program=program + ) + + def test_split_and_expand_performed(self): + """Test that preprocess returns the correct tapes when splitting and expanding + is needed.""" + + class NoMatOp(qml.operation.Operation): + """Dummy operation for expanding circuit.""" + + # pylint: disable=missing-function-docstring + num_wires = 1 + + # pylint: disable=arguments-renamed, invalid-overridden-method + @property + def has_matrix(self): + return False + + def decomposition(self): + return [qml.PauliX(self.wires), qml.PauliY(self.wires)] + + ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)] + # Need to specify grouping type to transform tape + measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))] + tapes = [ + qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]), + qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]), + ] + + dev = DefaultQubit() + config = qml.devices.ExecutionConfig(gradient_method="adjoint") + + program, new_config = dev.preprocess(config) + res_tapes, batch_fn = program(tapes) + expected_ops = [ + [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)], + [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)], + ] + + assert len(res_tapes) == 4 + for i, t in enumerate(res_tapes): + for op, expected_op in zip(t.operations, expected_ops[i % 2]): + qml.assert_equal(op, expected_op) + assert len(t.measurements) == 1 + if i < 2: + qml.assert_equal(t.measurements[0], measurements[0]) + else: + qml.assert_equal(t.measurements[0], measurements[1]) + + input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]) + assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])) + + assert new_config.grad_on_execution + assert new_config.use_device_gradient + + +def test_warning_if_not_device_batch_transform(): + """Test that a warning is raised if the users requests to not run device batch transform.""" + + # pylint: disable=too-few-public-methods + class CustomOp(qml.operation.Operator): + """Dummy operator.""" + + def decomposition(self): + return [qml.PauliX(self.wires[0])] + + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): + program, _ = dev.preprocess() + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + results = qml.execute( + [qs], dev, device_batch_transform=False, transform_program=program + ) + + assert len(results) == 1 + assert qml.math.allclose(results[0], -1) + + +@pytest.mark.parametrize("gradient_fn", (None, "backprop", qml.gradients.param_shift)) +def test_caching(gradient_fn): + """Test that cache execute returns the cached result if the same script is executed + multiple times, both in multiple times in a batch and in separate batches.""" + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + cache = {} + + with qml.Tracker(dev) as tracker: + results = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) + results2 = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) + + assert len(cache) == 1 + assert cache[qs.hash] == -1.0 + + assert list(results) == [-1.0, -1.0] + assert list(results2) == [-1.0, -1.0] + + assert tracker.totals["batches"] == 1 + assert tracker.totals["executions"] == 1 + assert cache[qs.hash] == -1.0 + + +class TestExecuteDeprecations: + """Class to test deprecation warnings in qml.execute. Warnings should be raised even if the default value is used.""" + + @pytest.mark.parametrize("expand_fn", (None, lambda qs: qs, "device")) + def test_expand_fn_is_deprecated(self, expand_fn): + """Test that expand_fn argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The expand_fn argument is deprecated" + ): + # None is a value used for expand_fn in the QNode + qml.execute([qs], dev, expand_fn=expand_fn) + + def test_max_expansion_is_deprecated(self): + """Test that max_expansion argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The max_expansion argument is deprecated" + ): + qml.execute([qs], dev, max_expansion=10) + + @pytest.mark.parametrize("override_shots", (False, 10)) + def test_override_shots_is_deprecated(self, override_shots): + """Test that override_shots argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([qs], dev, override_shots=override_shots) + + @pytest.mark.parametrize("device_batch_transform", (False, True)) + def test_device_batch_transform_is_deprecated(self, device_batch_transform): + """Test that device_batch_transform argument of qml.execute is deprecated.""" + # Need to use legacy device, otherwise another warning would be raised due to new Device interface + dev = qml.device("default.qubit.legacy", wires=1) + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) -def test_old_interface_no_device_jacobian_products(): - """Test that an error is always raised for the old device interface if device jacobian products are requested.""" - dev = qml.device("default.qubit.legacy", wires=2) - tape = qml.tape.QuantumScript([qml.RX(1.0, wires=0)], [qml.expval(qml.PauliZ(0))]) - with pytest.raises(qml.QuantumFunctionError): - qml.execute((tape,), dev, device_vjp=True) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + qml.execute([qs], dev, device_batch_transform=device_batch_transform) diff --git a/tests/interfaces/test_jacobian_products.py b/tests/interfaces/test_jacobian_products.py index 07f3ccb60ab..44510e71879 100644 --- a/tests/interfaces/test_jacobian_products.py +++ b/tests/interfaces/test_jacobian_products.py @@ -31,7 +31,8 @@ ) dev = qml.device("default.qubit") -dev_old = qml.device("default.qubit.legacy", wires=5) +with pytest.warns(qml.PennyLaneDeprecationWarning): + dev_old = qml.device("default.qubit.legacy", wires=5) dev_lightning = qml.device("lightning.qubit", wires=5) adjoint_config = qml.devices.ExecutionConfig(gradient_method="adjoint") dev_ps = ParamShiftDerivativesDevice() diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py index e9b7fdd8904..519c0daa028 100644 --- a/tests/interfaces/test_jax.py +++ b/tests/interfaces/test_jax.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,606 +11,456 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the JAX-Python interface""" +"""Jax specific tests for execute and default qubit 2.""" import numpy as np - -# pylint: disable=protected-access,too-few-public-methods import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import execute +from pennylane.devices import DefaultQubit from pennylane.gradients import param_shift -from pennylane.typing import TensorLike - -pytestmark = pytest.mark.jax +from pennylane.measurements import Shots jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") jax.config.update("jax_enable_x64", True) +pytestmark = pytest.mark.jax -class TestJaxExecuteUnitTests: - """Unit tests for jax execution""" - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - )[0] - - jax.grad(cost)(a, device=dev) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if an gradient transform - is used with grad_on_execution=True""" - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - grad_on_execution=True, - )[0] - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - jax.grad(cost)(a, device=dev) - - def test_unknown_interface(self): - """Test that an error is raised if the interface is unknown""" - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - interface="None", - )[0] - - with pytest.raises(ValueError, match="Unknown interface"): - cost(a, device=dev) - - def test_grad_on_execution(self, mocker): - """Test that grad_on_execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=2) - spy = mocker.spy(dev, "execute_and_gradients") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.expval(qml.PauliZ(0))], - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - ) - return execute( - [tape1, tape2, tape3], - dev, - gradient_fn="device", - gradient_kwargs={ - "method": "adjoint_jacobian", - "use_device_state": True, - }, - ) - - a = jax.numpy.array([0.1, 0.2]) - res = cost(a) - - x, y = a - assert np.allclose(res[0][0], np.cos(x) * np.cos(y)) - assert np.allclose(res[0][1], 1) - assert np.allclose(res[1], np.cos(0.5)) - assert np.allclose(res[2], np.cos(x) * np.cos(y)) - - # adjoint method only performs a single device execution per tape, but gets both result and gradient - assert dev.num_executions == 3 - spy.assert_not_called() - - g = jax.jacobian(cost)(a) - spy.assert_called() - expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y)) - assert qml.math.allclose(g[0][0], expected_g) - assert qml.math.allclose(g[0][1], np.zeros(2)) - assert qml.math.allclose(g[1], np.zeros(2)) - assert qml.math.allclose(g[2], expected_g) - - def test_no_grad_on_execution(self, mocker): - """Test that `grad_on_execution=False` uses the `device.execute_and_gradients`.""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - a = jax.numpy.array([0.1, 0.2]) - cost(a) +def test_jit_execution(): + """Test that qml.execute can be directly jitted.""" + dev = qml.device("default.qubit") - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() + tape = qml.tape.QuantumScript( + [qml.RX(jax.numpy.array(0.1), 0)], [qml.expval(qml.s_prod(2.0, qml.PauliZ(0)))] + ) - jax.grad(cost)(a) - spy_gradients.assert_called() + out = jax.jit(qml.execute, static_argnames=("device", "gradient_fn"))( + (tape,), device=dev, gradient_fn=qml.gradients.param_shift + ) + expected = 2.0 * jax.numpy.cos(jax.numpy.array(0.1)) + assert qml.math.allclose(out[0], expected) +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) + """Tests for caching behaviour""" - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cachesize=cachesize, - )[0] - - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cachesize=2) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 + @pytest.mark.skip("caching is not implemented for jax") + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + device = DefaultQubit() + params = jnp.arange(1, num_params + 1) / 10 - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + N = len(params) - def cost(a, cache): + def cost(x, cache): with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cache=cache, - )[0] - - custom_cache = {} - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_custom_cache_multiple(self, mocker): - """Test the use of a custom cache object with multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - - def cost(a, b, cache): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - res = execute( - [tape1, tape2], - dev, - gradient_fn=param_shift, - cache=cache, - ) - return res[0] - - custom_cache = {} - jax.grad(cost)(a, b, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) - def test_caching_param_shift(self, tol): - """Test that, when using parameter-shift transform, - caching produces the optimum number of evaluations.""" - dev = qml.device("default.qubit.legacy", wires=1) + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) + qml.CNOT(wires=[0, 1]) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cache=cache, + return qml.execute( + [tape], device, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 )[0] - # Without caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = jax.grad(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 - - # Check that calling the cost function again - # continues to evaluate the device (that is, the cache - # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 - assert np.allclose(grad1, grad2, atol=tol, rtol=0) - - # Check that calling the cost function again - # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 - assert not np.allclose(grad1, grad2, atol=tol, rtol=0) - - def test_caching_adjoint_backward(self): - """Test that caching produces the optimum number of adjoint evaluations - when no grad on execution.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) + # No caching: number of executions is not ideal + with qml.Tracker(device) as tracker: + hess1 = jax.jacobian(jax.grad(cost))(params, cache=False) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params + expected = jnp.array( + [ + [2 * jnp.cos(2 * x) * jnp.sin(y) ** 2, jnp.sin(2 * x) * jnp.sin(2 * y)], + [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs + + # Use caching: number of executions is ideal + + with qml.Tracker(device) as tracker2: + hess2 = jax.jacobian(jax.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert tracker2.totals["executions"] == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + +# add tests for lightning 2 when possible +# set rng for device when possible +no_shots = Shots(None) +shots_2_10k = Shots((10000, 10000)) +dev_def = DefaultQubit() +dev_ps = ParamShiftDerivativesDevice(seed=54353453) +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), # 0 + ({"gradient_fn": param_shift}, no_shots, dev_def), # 1 + ({"gradient_fn": "backprop"}, no_shots, dev_def), # 2 + ({"gradient_fn": "adjoint"}, no_shots, dev_def), # 3 + ({"gradient_fn": "adjoint", "device_vjp": True}, no_shots, dev_def), # 4 + ({"gradient_fn": "device"}, shots_2_10k, dev_ps), # 5 +] - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - # Without caching, 2 evaluations are required. - # 1 for the forward pass, and one per output dimension - # on the backward pass. - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 2 - - # With caching, also 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 3e-2 if shots else 1e-6 -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestJaxExecuteIntegration: """Test the jax interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) def cost(a, b): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) + return execute([tape1, tape2], device, **execute_kwargs) - return execute([tape1, tape2], dev, **execute_kwargs) + a = jnp.array(0.1) + b = np.array(0.2) + with device.tracker: + res = cost(a, b) - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - res = cost(a, b) + if execute_kwargs.get("gradient_fn", None) == "adjoint": + assert device.tracker.totals.get("execute_and_derivative_batches", 0) == 0 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" - a = jax.numpy.array(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + a = jnp.array(0.1) def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - res = jax.grad(cost)(a) - assert res.shape == () + res = jax.jacobian(cost)(a) + if not shots.has_partitioned_shots: + assert res.shape == () # pylint: disable=no-member # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -jnp.sin(a), atol=atol_for_shots(shots)) - def test_reusing_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + def test_jacobian(self, execute_kwargs, shots, device): + """Test jacobian calculation""" - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - assert tape.trainable_params == [0, 1] + a = jnp.array(0.1) + b = jnp.array(0.2) def cost(a, b): - # An explicit call to _update() is required here to update the - # trainable parameters in between tape executions. - # This is different from how the autograd interface works. - # Unless the update is issued, the validation check related to the - # number of provided parameters fails in the tape: (len(params) != - # required_length) and the tape produces incorrect results. - tape._update() - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return execute([new_tape], dev, **execute_kwargs)[0] + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return execute([tape], device, **execute_kwargs)[0] - jac_fn = jax.grad(cost) - jac = jac_fn(a, b) + res = cost(a, b) + expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)] + if shots.has_partitioned_shots: + assert np.allclose(res[0], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected, atol=2 * atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + g = jax.jacobian(cost, argnums=[0, 1])(a, b) + assert isinstance(g, tuple) and len(g) == 2 + + expected = ([-jnp.sin(a), jnp.sin(a) * jnp.sin(b)], [0, -jnp.cos(a) * jnp.cos(b)]) + + if shots.has_partitioned_shots: + for i in (0, 1): + assert np.allclose(g[i][0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(g[0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) + + def test_tape_no_parameters(self, execute_kwargs, shots, device): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" - a = jax.numpy.array(0.54) - b = jax.numpy.array(0.8) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [np.cos(2 * a)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + tape2 = qml.tape.QuantumScript( + [qml.RY(jnp.array(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - jac_fn = jax.grad(lambda a, b: cost(2 * a, b)) - jac = jac_fn(a, b) - expected = -2 * np.sin(2 * a) - assert np.allclose(jac, expected, atol=tol, rtol=0) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - def test_grad_with_different_grad_on_execution(self, execute_kwargs): - """Test jax grad for adjoint diff method with different execution kwargs.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + tape4 = qml.tape.QuantumScript( + [qml.RY(jnp.array(0.5), 0)], [qml.probs(wires=[0, 1])], shots=shots + ) + res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + res = jax.tree_util.tree_leaves(res) + out = sum(jnp.hstack(res)) + if shots.has_partitioned_shots: + out = out / shots.num_copies + return out - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) + params = jnp.array([0.1, 0.2]) - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] - return res + x, y = params - results = jax.grad(cost)(params, cache=None) - for r, e in zip(results, expected_results): - assert jax.numpy.allclose(r, e, atol=1e-7) + res = cost(params) + expected = 2 + jnp.cos(0.5) + jnp.cos(x) * jnp.cos(y) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing_single_tape(self, execute_kwargs): - """Test classical processing within the quantum tape for a single tape""" - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - c = jax.numpy.array(0.3) + grad = jax.grad(cost)(params) + expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)] + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - def cost(a, b, c, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + # pylint: disable=too-many-statements + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - return execute([tape], device, **execute_kwargs)[0] + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - dev = qml.device("default.qubit.legacy", wires=2) - res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev) - assert len(res) == 3 + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], [qml.expval(qml.PauliZ(0))], shots=shots + ) - def test_classical_processing_multiple_tapes(self, execute_kwargs): - """Test classical processing within the quantum tape for multiple - tapes""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = qml.execute([tape1, tape2, tape3], device, **execute_kwargs) + leaves = jax.tree_util.tree_leaves(res) + return jnp.hstack(leaves) + + params = jnp.array([0.1, 0.2]) + x, y = params + + res = cost(params) + assert isinstance(res, jax.Array) + assert res.shape == (4 * shots.num_copies,) if shots.has_partitioned_shots else (4,) + + if shots.has_partitioned_shots: + assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[2], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[3], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[4], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[5], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[6], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[7], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + + jac = jax.jacobian(cost)(params) + assert isinstance(jac, jnp.ndarray) + assert ( + jac.shape == (8, 2) if shots.has_partitioned_shots else (4, 2) + ) # pylint: disable=no-member + + if shots.has_partitioned_shots: + assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) + assert np.allclose(jac[3:5], 0, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) + + d1 = -jnp.sin(x) * jnp.cos(y) + if shots.has_partitioned_shots: + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) + + d2 = -jnp.cos(x) * jnp.sin(y) + if shots.has_partitioned_shots: + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + if execute_kwargs["gradient_fn"] == param_shift: + pytest.skip("Basic QNode execution wipes out trainable params with param-shift") - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.RY(x[0], wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) + a = jnp.array(0.1) + b = jnp.array(0.2) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.Hadamard(0) - qml.CRX(2 * x[0] * x[1], wires=[0, 1]) - qml.RX(2 * x[1], wires=[1]) - qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + shots=shots, + ) + assert tape.trainable_params == [0, 1] - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) - return result[0] + result[1] - 7 * result[1] + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return jnp.hstack(execute([new_tape], device, **execute_kwargs)[0]) - res = jax.grad(cost_fn)(params) - assert res.shape == (2,) + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + jac = jac_fn(a, b) - def test_multiple_tapes_output(self, execute_kwargs): - """Test the output types for the execution of multiple quantum tapes""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + a = jnp.array(0.54) + b = jnp.array(0.8) - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.RY(x[0], wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [jnp.cos(2 * a), -jnp.cos(2 * a) * jnp.sin(b)] + if shots.has_partitioned_shots: + assert np.allclose(res2[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res2[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(lambda a, b: cost(2 * a, b), argnums=[0, 1]) + jac = jac_fn(a, b) + expected = ( + [-2 * jnp.sin(2 * a), 2 * jnp.sin(2 * a) * jnp.sin(b)], + [0, -jnp.cos(2 * a) * jnp.cos(b)], + ) + assert isinstance(jac, tuple) and len(jac) == 2 + if shots.has_partitioned_shots: + for offset in (0, 2): + assert np.allclose(jac[0][0 + offset], expected[0][0], atol=atol_for_shots(shots)) + assert np.allclose(jac[0][1 + offset], expected[0][1], atol=atol_for_shots(shots)) + assert np.allclose(jac[1][0 + offset], expected[1][0], atol=atol_for_shots(shots)) + assert np.allclose(jac[1][1 + offset], expected[1][1], atol=atol_for_shots(shots)) + else: + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) + + def test_classical_processing(self, execute_kwargs, shots, device): + """Test classical processing within the quantum tape""" + a = jnp.array(0.1) + b = jnp.array(0.2) + c = jnp.array(0.3) + + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + jnp.sin(a), wires=0), + ] - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.Hadamard(0) - qml.CRX(2 * x[0] * x[1], wires=[0, 1]) - qml.RX(2 * x[1], wires=[1]) - qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - tape2 = qml.tape.QuantumScript.from_queue(q2) + res = jax.jacobian(cost, argnums=[0, 2])(a, b, c) - return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - res = cost_fn(params) - assert isinstance(res, TensorLike) - assert all(isinstance(r, jax.numpy.ndarray) for r in res) - assert all(r.shape == () for r in res) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - def test_matrix_parameter(self, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the jax interface works correctly with a matrix parameter""" - a = jax.numpy.array(0.1) - U = jax.numpy.array([[0, 1], [1, 0]]) - - def cost(a, U, device): - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + U = jnp.array([[0, 1], [1, 0]]) + a = jnp.array(0.1) - tape = qml.tape.QuantumScript.from_queue(q) - tape.trainable_params = [0] + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, U, device=dev) - assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + res = cost(a, U) + assert np.allclose(res, -jnp.cos(a), atol=atol_for_shots(shots), rtol=0) - jac_fn = jax.grad(cost, argnums=0) - res = jac_fn(a, U, device=dev) - assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + jac_fn = jax.jacobian(cost) + jac = jac_fn(a, U) + if not shots.has_partitioned_shots: + assert isinstance(jac, jnp.ndarray) + assert np.allclose(jac, jnp.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): + def test_differentiable_expand(self, execute_kwargs, device, shots): """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -619,258 +469,349 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - def cost_fn(a, p, device): - with qml.queuing.AnnotatedQueue() as q_tape: - qml.RX(a, wires=0) - U3(*p, wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q_tape) - tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) - return execute([tape], device, **execute_kwargs)[0] + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], + [qml.expval(qml.PauliX(0))], + shots=shots, + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + conf = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=conf) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - a = jax.numpy.array(0.1) - p = jax.numpy.array([0.1, 0.2, 0.3]) + a = jnp.array(0.1) + p = jnp.array([0.1, 0.2, 0.3]) - dev = qml.device("default.qubit.legacy", wires=1) - res = cost_fn(a, p, device=dev) - expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( - np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + res = cost_fn(a, p) + expected = jnp.cos(a) * jnp.cos(p[1]) * jnp.sin(p[0]) + jnp.sin(a) * ( + jnp.cos(p[2]) * jnp.sin(p[1]) + jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - jac_fn = jax.grad(cost_fn, argnums=1) - res = jac_fn(a, p, device=dev) - expected = jax.numpy.array( + jac_fn = jax.jacobian(cost_fn, argnums=[1]) + res = jac_fn(a, p) + expected = jnp.array( [ - np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), - np.cos(p[1]) * np.cos(p[2]) * np.sin(a) - - np.sin(p[1]) - * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), - np.sin(a) - * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + jnp.cos(p[1]) + * (jnp.cos(a) * jnp.cos(p[0]) - jnp.sin(a) * jnp.sin(p[0]) * jnp.sin(p[2])), + jnp.cos(p[1]) * jnp.cos(p[2]) * jnp.sin(a) + - jnp.sin(p[1]) + * (jnp.cos(a) * jnp.sin(p[0]) + jnp.cos(p[0]) * jnp.sin(a) * jnp.sin(p[2])), + jnp.sin(a) + * (jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.cos(p[2]) - jnp.sin(p[1]) * jnp.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_independent_expval(self, execute_kwargs): - """Tests computing an expectation value that is independent of trainable - parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = jax.grad(cost)(params, cache=None) - assert res.shape == (3,) + def test_probability_differentiation(self, execute_kwargs, device, shots): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return jnp.hstack(execute([tape], device, **execute_kwargs)[0]) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) -class TestVectorValued: - """Test vector-valued jacobian returns for the JAX Python interface.""" - - def test_multiple_expvals(self, execute_kwargs): - """Tests computing multiple expectation values in a tape.""" - - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] - - res = jax.jacobian(cost)(params, cache=None) + x = jnp.array(0.543) + y = jnp.array(-0.654) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert res[0].shape == (3,) - assert isinstance(res[0], jax.numpy.ndarray) - - assert res[1].shape == (3,) - assert isinstance(res[1], jax.numpy.ndarray) - - def test_multiple_expvals_single_par(self, execute_kwargs): - """Tests computing multiple expectation values in a tape with a single - trainable parameter.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] - - res = jax.jacobian(cost)(params, cache=None) - - assert isinstance(res, tuple) - - assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (1,) - - assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (1,) + res = cost(x, y) + expected = jnp.array( + [ + [ + jnp.cos(x / 2) ** 2, + jnp.sin(x / 2) ** 2, + (1 + jnp.cos(x) * jnp.cos(y)) / 2, + (1 - jnp.cos(x) * jnp.cos(y)) / 2, + ], + ] + ) + if shots.has_partitioned_shots: + assert np.allclose(res[:, 0:2].flatten(), expected, atol=atol_for_shots(shots)) + assert np.allclose(res[:, 2:].flatten(), expected, atol=atol_for_shots(shots)) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + res = jac_fn(x, y) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,) + assert res[1].shape == (2, 4) if shots.has_partitioned_shots else (4,) + + expected = ( + jnp.array( + [ + [ + -jnp.sin(x) / 2, + jnp.sin(x) / 2, + -jnp.sin(x) * jnp.cos(y) / 2, + jnp.sin(x) * jnp.cos(y) / 2, + ], + ] + ), + jnp.array( + [ + [0, 0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2], + ] + ), + ) - def test_multi_tape_fwd(self, execute_kwargs): - """Test the forward evaluation of a cost function that uses the output - of multiple tapes that be vector-valued.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + if shots.has_partitioned_shots: + assert np.allclose(res[0][:, 0:2].flatten(), expected[0], atol=atol_for_shots(shots)) + assert np.allclose(res[0][:, 2:].flatten(), expected[0], atol=atol_for_shots(shots)) + assert np.allclose(res[1][:, :2].flatten(), expected[1], atol=atol_for_shots(shots)) + assert np.allclose(res[1][:, 2:].flatten(), expected[1], atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, device, shots): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + res = qml.execute([tape], device, **execute_kwargs)[0] + return jnp.hstack(jax.tree_util.tree_leaves(res)) + + x = jnp.array(0.543) + y = jnp.array(-0.654) + + res = cost(x, y) + expected = jnp.array( + [jnp.cos(x), (1 + jnp.cos(x) * jnp.cos(y)) / 2, (1 - jnp.cos(x) * jnp.cos(y)) / 2] + ) + if shots.has_partitioned_shots: + assert np.allclose(res[:3], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[3:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + res = jac_fn(x, y) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) + assert res[1].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) + + expected = ( + jnp.array([-jnp.sin(x), -jnp.sin(x) * jnp.cos(y) / 2, jnp.sin(x) * jnp.cos(y) / 2]), + jnp.array([0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2]), + ) + if shots.has_partitioned_shots: + assert np.allclose(res[0][:3], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[0][3:], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1][:3], expected[1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1][3:], expected[1], atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + +class TestHigherOrderDerivatives: + """Test that the jax execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + jnp.array([0.543, -0.654]), + jnp.array([0, -0.654]), + jnp.array([-2.0, 0]), + ], + ) + def test_parameter_shift_hessian(self, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using jax, yielding second derivatives.""" + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.expval(qml.PauliY(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[1], wires=[0]) - qml.RX(x[1], wires=[0]) - qml.RX(-x[1], wires=[0]) - qml.expval(qml.PauliY(0)) - qml.expval(qml.PauliY(1)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) + + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] - expected = -jax.numpy.sin(params[0]) + -jax.numpy.sin(params[1]) res = cost_fn(params) - assert jax.numpy.allclose(expected, res) - - def test_multi_tape_jacobian(self, execute_kwargs): - """Test the jacobian computation with multiple tapes.""" + x, y = params + expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - def cost(x, y, device, interface, ek): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + res = jax.grad(cost_fn)(params) + expected = jnp.array( + [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + res = jax.jacobian(jax.grad(cost_fn))(params) + expected = jnp.array( + [ + [-jnp.cos(2 * x) * jnp.cos(2 * y), jnp.sin(2 * x) * jnp.sin(2 * y)], + [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], device, **ek, interface=interface) + def test_max_diff(self, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = DefaultQubit() + params = jnp.array([0.543, -0.654]) - dev = qml.device("default.qubit.legacy", wires=2) - x = jax.numpy.array(0.543) - y = jax.numpy.array(-0.654) + def cost_fn(x): + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - x_ = np.array(0.543) - y_ = np.array(-0.654) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) - exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] - assert np.allclose(exec_jax, exec_autograd) + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - res = jax.jacobian(cost, argnums=(0, 1))( - x, y, dev, interface="jax-python", ek=execute_kwargs + res = jax.grad(cost_fn)(params) + expected = jnp.array( + [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] ) + assert np.allclose(res, expected, atol=tol, rtol=0) - import autograd.numpy as anp - - def cost_stack(x, y, device, interface, ek): - return anp.hstack(cost(x, y, device, interface, ek)) + res = jax.jacobian(jax.grad(cost_fn))(params) + expected = jnp.zeros([2, 2]) + assert np.allclose(res, expected, atol=tol, rtol=0) - exp = qml.jacobian(cost_stack, argnum=(0, 1))( - x_, y_, dev, interface="autograd", ek=execute_kwargs - ) - res_0 = jax.numpy.array([res[0][0][0], res[0][1][0], res[1][0][0], res[1][1][0]]) - res_1 = jax.numpy.array([res[0][0][1], res[0][1][1], res[1][0][1], res[1][1][1]]) - assert np.allclose(res_0, exp[0]) - assert np.allclose(res_1, exp[1]) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" - def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): - """Test the jacobian computation with multiple tapes with probability - and expectation value computations.""" + @pytest.fixture + def cost_fn(self, execute_kwargs, shots, device): + """Cost function for gradient tests""" - adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" - if adjoint: - pytest.skip("The adjoint diff method doesn't support probabilities.") + def _cost_fn(weights, coeffs1, coeffs2): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() - def cost(x, y, device, interface, ek): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], device, **ek, interface=interface) - - dev = qml.device("default.qubit.legacy", wires=2) - x = jax.numpy.array(0.543) - y = jax.numpy.array(-0.654) - - x_ = np.array(0.543) - y_ = np.array(-0.654) - - exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) - exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) - - assert all( - np.allclose(exec_jax[i][j], exec_autograd[i][j]) for i in range(2) for j in range(2) - ) - - res = jax.jacobian(cost, argnums=(0, 1))( - x, y, dev, interface="jax-python", ek=execute_kwargs + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + res = execute([tape], device, **execute_kwargs)[0] + if shots.has_partitioned_shots: + return jnp.hstack(res[0] + res[1]) + return jnp.hstack(res) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * jnp.sin(x) * jnp.sin(y) + jnp.cos(x) * (a + b * jnp.sin(y)), d * jnp.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return jnp.array( + [ + [ + -c * jnp.cos(x) * jnp.sin(y) - jnp.sin(x) * (a + b * jnp.sin(y)), + b * jnp.cos(x) * jnp.cos(y) - c * jnp.cos(y) * jnp.sin(x), + jnp.cos(x), + jnp.cos(x) * jnp.sin(y), + -(jnp.sin(x) * jnp.sin(y)), + 0, + ], + [-d * jnp.sin(x), 0, 0, 0, 0, jnp.cos(x)], + ] ) - assert isinstance(res, TensorLike) - assert len(res) == 2 - - for r, exp_shape in zip(res, [(), (2,)]): - assert isinstance(r, tuple) - assert len(r) == 2 - assert len(r[0]) == 2 - assert isinstance(r[0][0], jax.numpy.ndarray) - assert r[0][0].shape == exp_shape - assert isinstance(r[0][1], jax.numpy.ndarray) - assert r[0][1].shape == exp_shape - assert len(r[1]) == 2 - assert isinstance(r[1][0], jax.numpy.ndarray) - assert r[1][0].shape == exp_shape - assert isinstance(r[1][1], jax.numpy.ndarray) - assert r[1][1].shape == exp_shape + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + coeffs1 = jnp.array([0.1, 0.2, 0.3]) + coeffs2 = jnp.array([0.7]) + weights = jnp.array([0.4, 0.5]) + + res = cost_fn(weights, coeffs1, coeffs2) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + res = jax.jacobian(cost_fn)(weights, coeffs1, coeffs2) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + if shots.has_partitioned_shots: + assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") + + coeffs1 = jnp.array([0.1, 0.2, 0.3]) + coeffs2 = jnp.array([0.7]) + weights = jnp.array([0.4, 0.5]) + + res = cost_fn(weights, coeffs1, coeffs2) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + res = jnp.hstack(jax.jacobian(cost_fn, argnums=[0, 1, 2])(weights, coeffs1, coeffs2)) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index c3fd7df5d1a..4fe3d4946d2 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -14,7 +14,7 @@ """Unit tests for the JAX-JIT interface""" import numpy as np -# pylint: disable=protected-access,too-few-public-methods +# pylint: disable=protected-access,too-few-public-methods,unnecessary-lambda import pytest import pennylane as qml @@ -37,7 +37,7 @@ def test_jacobian_options(self, mocker): a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -64,7 +64,7 @@ def test_incorrect_gradients_on_execution(self): is used with grad_on_execution=True""" a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -90,7 +90,7 @@ def test_unknown_interface(self): """Test that an error is raised if the interface is unknown""" a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -112,8 +112,8 @@ def cost(a, device): def test_grad_on_execution(self, mocker): """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") + dev = qml.device("default.qubit", wires=1) + spy = mocker.spy(dev, "execute_and_compute_derivatives") def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -134,22 +134,22 @@ def cost(a): )[0] a = jax.numpy.array([0.1, 0.2]) - jax.jit(cost)(a) + + with qml.Tracker(dev) as tracker: + jax.jit(cost)(a) # adjoint method only performs a single device execution # gradients are not requested when we only want the results - assert dev.num_executions == 1 spy.assert_not_called() + assert tracker.totals["executions"] == 1 # when the jacobian is requested, we always calculate it at the same time as the results jax.grad(jax.jit(cost))(a) spy.assert_called() - def test_no_gradients_on_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + def test_no_gradients_on_execution(self): + """Test that no grad on execution uses the `device.execute` and `device.compute_derivatives` pathway""" + dev = qml.device("default.qubit", wires=1) def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -168,14 +168,16 @@ def cost(a): )[0] a = jax.numpy.array([0.1, 0.2]) - jax.jit(cost)(a) - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() + with qml.Tracker(dev) as tracker: + jax.jit(cost)(a) - jax.grad(jax.jit(cost))(a) - spy_gradients.assert_called() + assert tracker.totals["executions"] == 1 + assert "derivatives" not in tracker.totals + + with qml.Tracker(dev) as tracker: + jax.grad(jax.jit(cost))(a) + assert tracker.totals["derivatives"] == 1 class TestCaching: @@ -183,7 +185,7 @@ class TestCaching: def test_cache_maxsize(self, mocker): """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): @@ -211,7 +213,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): @@ -238,7 +240,7 @@ def cost(a, cache): def test_custom_cache_multiple(self, mocker): """Test the use of a custom cache object with multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = jax.numpy.array(0.1) @@ -276,7 +278,7 @@ def cost(a, b, cache): def test_caching_param_shift(self, tol): """Test that, when using parameter-shift transform, caching produces the optimum number of evaluations.""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -296,33 +298,37 @@ def cost(a, cache): # Without caching, 5 evaluations are required to compute # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 5 + with qml.Tracker(dev) as tracker: + jax.grad(cost)(params, cache=None) + + assert tracker.totals["executions"] == 5 # With caching, 5 evaluations are required to compute # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = jax.grad(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 + with qml.Tracker(dev) as tracker: + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert tracker.totals["executions"] == 5 # Check that calling the cost function again # continues to evaluate the device (that is, the cache # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 + with qml.Tracker(dev) as tracker: + grad2 = jac_fn(params, cache=True) + assert tracker.totals["executions"] == 5 assert np.allclose(grad1, grad2, atol=tol, rtol=0) # Check that calling the cost function again # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 + with qml.Tracker(dev) as tracker: + grad2 = jac_fn(2 * params, cache=True) + assert tracker.totals["executions"] == 5 assert not np.allclose(grad1, grad2, atol=tol, rtol=0) def test_caching_adjoint_backward(self): """Test that caching produces the optimum number of adjoint evaluations when mode=backward""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -346,15 +352,18 @@ def cost(a, cache): # Without caching, 2 evaluations are required. # 1 for the forward pass, and one per output dimension # on the backward pass. - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 2 + with qml.Tracker(dev) as tracker: + jax.grad(cost)(params, cache=None) + assert tracker.totals["executions"] == 1 + assert tracker.totals["derivatives"] == 1 # With caching, also 2 evaluations are required. One # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 + with qml.Tracker(dev) as tracker: + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert tracker.totals["executions"] == 1 + assert tracker.totals["derivatives"] == 1 execute_kwargs_integration = [ @@ -379,7 +388,7 @@ class TestJaxExecuteIntegration: def test_execution(self, execute_kwargs): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, b): with qml.queuing.AnnotatedQueue() as q1: @@ -409,7 +418,7 @@ def cost(a, b): def test_scalar_jacobian(self, execute_kwargs, tol): """Test scalar jacobian calculation""" a = jax.numpy.array(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -432,7 +441,7 @@ def cost(a): tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(dev.execute(tapes)) assert expected.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) @@ -442,7 +451,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) with qml.queuing.AnnotatedQueue() as q: qml.RY(a, wires=0) @@ -484,7 +493,7 @@ def cost(a, b): def test_grad_with_backward_mode(self, execute_kwargs): """Test jax grad for adjoint diff method in backward mode""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) @@ -523,14 +532,14 @@ def cost(a, b, c, device): return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) res = jax.jit(jax.grad(cost, argnums=(0, 1, 2)), static_argnums=3)(a, b, c, device=dev) assert len(res) == 3 def test_classical_processing_multiple_tapes(self, execute_kwargs): """Test classical processing within the quantum tape for multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.3, 0.2]) def cost_fn(x): @@ -558,7 +567,7 @@ def cost_fn(x): def test_multiple_tapes_output(self, execute_kwargs): """Test the output types for the execution of multiple quantum tapes""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.3, 0.2]) def cost_fn(x): @@ -602,7 +611,7 @@ def cost(a, U, device): tape.trainable_params = [0] return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) res = jax.jit(cost, static_argnums=2)(a, U, device=dev) assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) @@ -627,13 +636,15 @@ def cost_fn(a, p, device): qscript = qml.tape.QuantumScript( [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] ) - qscript = qscript.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + qscript = qscript.expand( + stop_at=lambda obj: qml.devices.default_qubit.stopping_condition(obj) + ) return execute([qscript], device, **execute_kwargs)[0] a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev) expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) @@ -657,7 +668,7 @@ def cost_fn(a, p, device): def test_independent_expval(self, execute_kwargs): """Tests computing an expectation value that is independent of trainable parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -694,7 +705,7 @@ def test_shapes(self, execute_kwargs, ret_type, shape, expected_type): if adjoint: pytest.skip("The adjoint diff method doesn't support probabilities.") - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -722,7 +733,7 @@ def cost(a, cache): def test_independent_expval(self, execute_kwargs): """Tests computing an expectation value that is independent of trainable parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -754,7 +765,7 @@ def cost(a, cache): def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim, expected_type): """Tests the shape of vector-valued QNode results.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) grad_meth = ( execute_kwargs["gradient_kwargs"]["method"] @@ -793,7 +804,7 @@ def cost(a, cache): def test_qnode_sample(self, execute_kwargs): """Tests computing multiple expectation values in a tape.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) + dev = qml.device("default.qubit", wires=2, shots=10) params = jax.numpy.array([0.1, 0.2, 0.3]) grad_meth = ( @@ -817,11 +828,12 @@ def cost(a, cache): return res res = jax.jit(cost, static_argnums=1)(params, cache=None) - assert res.shape == (dev.shots,) + + assert res.shape == (dev.shots.total_shots,) def test_multiple_expvals_grad(self, execute_kwargs): """Tests computing multiple expectation values in a tape.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -868,7 +880,7 @@ def cost(x, y, device, interface, ek): return qml.execute([tape1, tape2], device, **ek, interface=interface)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) diff --git a/tests/interfaces/test_jax_jit_qnode.py b/tests/interfaces/test_jax_jit_qnode.py index f109ddc0a4c..1c0b8ad1f59 100644 --- a/tests/interfaces/test_jax_jit_qnode.py +++ b/tests/interfaces/test_jax_jit_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,21 +12,35 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-JIT interface with a QNode""" +import copy + # pylint: disable=too-many-arguments,too-few-public-methods +from functools import partial + import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", True, True], + [ParamShiftDerivativesDevice(), "device", False, True], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "parameter-shift", False, False], ] interface_and_qubit_device_and_diff_method = [ ["auto"] + inner_list for inner_list in qubit_device_and_diff_method @@ -43,33 +57,33 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, interface, device_vjp + ): """Test execution works with the interface""" - if diff_method == "backprop": - pytest.skip("Test does not support backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = np.array(0.1, requires_grad=True) + a = jax.numpy.array(0.1) jax.jit(circuit)(a) assert circuit.interface == interface @@ -82,7 +96,9 @@ def circuit(a): assert isinstance(grad, jax.Array) assert grad.shape == () - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -91,13 +107,12 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, i a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=2) - @qnode( dev, interface=interface, diff_method="parameter-shift", grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -130,21 +145,21 @@ def circuit(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -159,21 +174,24 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the jax interface works correctly with a matrix parameter""" - U = jax.numpy.array([[0, 1], [1, 0]]) - a = jax.numpy.array(0.1) - num_wires = 2 + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + U = jax.numpy.array([[0, 1], [1, 0]]) + a = jax.numpy.array(0.1) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -186,16 +204,19 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that operation and nested tape expansion is differentiable""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs = { - "sampler_rng": SEED_FOR_SPSA, - "num_directions": 10, - } + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA class U3(qml.U3): @@ -207,13 +228,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -222,6 +236,7 @@ def decomposition(self): diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, p): @@ -248,15 +263,12 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = np.array([0.1, 0.2], requires_grad=True) - dev = qml.device(dev_name, wires=1) - @qnode( dev, interface=interface, @@ -264,6 +276,7 @@ def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interf h=1e-8, approx_order=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -278,36 +291,37 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_expval( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation""" - + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -349,29 +363,30 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_jacobian_no_evaluate( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} + if diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -412,22 +427,21 @@ def circuit(a, b): assert r.shape == () assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_single_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with a single prob output""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -436,6 +450,7 @@ def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, inter diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -465,22 +480,21 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_multi_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -489,6 +503,7 @@ def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -551,22 +566,21 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -575,6 +589,7 @@ def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, inter diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -628,23 +643,19 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - tol = TOL_FOR_SPSA + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -653,6 +664,7 @@ def test_diff_expval_probs_sub_argnums( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -685,19 +697,21 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "hadamard": + if diff_method == "hadamard": pytest.skip("Hadamard does not support var") elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -706,6 +720,7 @@ def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interfac diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -766,8 +781,8 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test jax device works with diff_method=None.""" - dev = qml.device("default.qubit.jax", wires=1, shots=10) + """Test device works with diff_method=None.""" + dev = DefaultQubit() @jax.jit @qml.qnode(dev, diff_method=None, interface=interface) @@ -777,52 +792,40 @@ def circuit(x): assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - def test_changing_shots(self, interface, mocker, tol): + @pytest.mark.skip("jax.jit does not work with sample") + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # no additional calls + assert res.shape == (100, 2) # pylint:disable=comparison-with-callable def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - # TODO: jit when https://github.com/PennyLaneAI/pennylane/issues/3474 is resolved - res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) - assert dev.shots == 1 + jit_cost_fn = jax.jit(cost_fn, static_argnames=["shots"]) + res = jax.grad(jit_cost_fn, argnums=[0, 1])(a, b, shots=30000) expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) @@ -830,44 +833,79 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - # We're choosing interface="jax" such that backprop can be used in the - # test later - @qnode(dev, interface="jax") + @qnode(DefaultQubit(), interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + cost_fn(a, b) + assert cost_fn.gradient_fn == "backprop" # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) + def test_shot_vectors_single_measurements(self, interface, shots): + """Test jax-jit can work with shot vectors.""" + + dev = qml.device("default.qubit", shots=shots, seed=4747) + + @jax.jit + @qml.qnode(dev, interface=interface, diff_method="parameter-shift") + def circuit(x): + qml.RX(x, wires=0) + return qml.var(qml.PauliZ(0)) + + res = circuit(0.5) + expected = 1 - np.cos(0.5) ** 2 + assert qml.math.allclose(res[0], expected, atol=1e-2) + assert qml.math.allclose(res[1], expected, atol=3e-2) + + g = jax.jacobian(circuit)(0.5) + + expected_g = 2 * np.cos(0.5) * np.sin(0.5) + assert qml.math.allclose(g[0], expected_g, atol=2e-2) + assert qml.math.allclose(g[1], expected_g, atol=2e-2) + + @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) + def test_shot_vectors_multiple_measurements(self, interface, shots): + """Test jax-jit can work with shot vectors.""" + + dev = qml.device("default.qubit", shots=shots, seed=987548) + + @jax.jit + @qml.qnode(dev, interface=interface, diff_method="parameter-shift") + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=0) + + res = circuit(0.5) + assert qml.math.allclose(res[0][0], np.cos(0.5), atol=5e-3) + assert qml.math.allclose(res[1][0], np.cos(0.5), atol=5e-3) + expected_probs = np.array([np.cos(0.25) ** 2, np.sin(0.25) ** 2]) + assert qml.math.allclose(res[0][1], expected_probs, atol=5e-3) + assert qml.math.allclose(res[1][1], expected_probs, atol=5e-3) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): + def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test sampling works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -875,17 +913,19 @@ def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) + return qml.sample(qml.Z(0)), qml.sample(qml.s_prod(2, qml.X(0) @ qml.Y(1))) - res = jax.jit(circuit)() + res = jax.jit(circuit, static_argnames="shots")(shots=10) assert isinstance(res, tuple) @@ -894,7 +934,7 @@ def circuit(): assert isinstance(res[1], jax.Array) assert res[1].shape == (10,) - def test_counts(self, dev_name, diff_method, grad_on_execution, interface): + def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test counts works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -902,10 +942,12 @@ def test_counts(self, dev_name, diff_method, grad_on_execution, interface): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): qml.Hadamard(wires=[0]) @@ -916,9 +958,9 @@ def circuit(): with pytest.raises( NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): - jax.jit(circuit)() + jax.jit(circuit, static_argnames="shots")(shots=10) else: - res = jax.jit(circuit)() + res = jax.jit(circuit, static_argnames="shots")(shots=10) assert isinstance(res, tuple) @@ -927,28 +969,32 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution, interface): + def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) @@ -974,35 +1020,76 @@ def cost(weights): assert len(res) == 2 + def test_postselection_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" + + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + elif dev.name == "lightning.qubit": + pytest.xfail("lightning qubit does not support postselection.") + + @qml.qnode( + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + phi = jax.numpy.array(1.23) + theta = jax.numpy.array(4.56) + + assert np.allclose(jax.jit(circuit)(phi, theta), jax.jit(expected_circuit)(theta)) + + gradient = jax.jit(jax.grad(circuit, argnums=[0, 1]))(phi, theta) + exp_theta_grad = jax.jit(jax.grad(expected_circuit))(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) + @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test second derivative calculation of a scalar-valued QNode""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not second derivative.") + if diff_method in {"adjoint", "device"}: + pytest.skip("Adjoint does not support second derivatives.") elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "num_directions": 10} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 + gradient_kwargs["h"] = H_FOR_SPSA tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1033,31 +1120,25 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { "h": H_FOR_SPSA, - "num_directions": 20, + "num_directions": 40, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1091,10 +1172,12 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1104,18 +1187,12 @@ def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, i } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1159,11 +1236,11 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, interface, grad_on_execution, tol + self, dev, diff_method, interface, device_vjp, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1173,18 +1250,12 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1231,11 +1302,11 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1245,19 +1316,13 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -1282,6 +1347,7 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) + hess = jax.jit(jax.jacobian(jac_fn, argnums=[0, 1]))(a, b) expected_hess = np.array( @@ -1301,23 +1367,24 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "lightning.qubit" and diff_method == "adjoint": + pytest.xfail("lightning.qubit does not support adjoint with the state.") x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) + if not dev.wires: + dev = copy.copy(dev) + dev._wires = qml.wires.Wires([0, 1]) # pylint:disable=protected-access @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1327,7 +1394,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") # pylint:disable=no-member + assert res.dtype is np.dtype("complex128") probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1341,9 +1408,13 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("Adjoint does not support projectors") elif diff_method == "hadamard": @@ -1352,7 +1423,6 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1361,6 +1431,7 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1383,97 +1454,9 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) -# TODO: Add CV test when return types and custom diff are compatible -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -@pytest.mark.parametrize("interface", ["auto", "jax-jit", "jax"]) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = 0.543 - phi = -0.654 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(r, phi) - expected = np.array( - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - n = 0.12 - a = 0.765 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(n, a) - expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -# TODO: add support for fwd grad_on_execution to JAX-JIT -@pytest.mark.parametrize("interface", ["auto", "jax-jit"]) -def test_adjoint_reuse_device_state(mocker, interface): - """Tests that the jax interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - jax.grad(circ)(1.0) - assert circ.device.num_executions == 1 - - spy.assert_called_with(mocker.ANY, use_device_state=True) - - @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1481,20 +1464,13 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - class PhaseShift(qml.PhaseShift): grad_method = None @@ -1507,6 +1483,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1517,16 +1494,19 @@ def circuit(x, y): x = jax.numpy.array(0.5) y = jax.numpy.array(0.7) circuit(x, y) + jax.grad(circuit, argnums=[0])(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol + self, dev, diff_method, grad_on_execution, max_diff, device_vjp, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} + if dev.name == "param_shift.qubit": + pytest.xfail("gradients transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "hadamard": @@ -1539,16 +1519,17 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] + @jax.jit @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1592,14 +1573,16 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method in ("adjoint", "backprop", "finite-diff"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "hadamard": @@ -1608,8 +1591,6 @@ def test_hamiltonian_expansion_finite_shots( gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "h": H_FOR_SPSA, "num_directions": 20} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1618,6 +1599,7 @@ def test_hamiltonian_expansion_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1633,13 +1615,12 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1664,27 +1645,21 @@ def circuit(data, weights, coeffs): # assert np.allclose(grad2_w_c, expected, atol=tol) def test_vmap_compared_param_broadcasting( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not yet support Hamiltonians") - elif diff_method == "hadamard": - pytest.skip("The Hadamard method does not yet support Hamiltonians") - - if diff_method == "backprop": - pytest.skip( - "The backprop method does not yet support parameter-broadcasting with Hamiltonians" - ) - phys_qubits = 2 - if diff_method == "hadamard": - phys_qubits = 3 n_configs = 5 pars_q = np.random.rand(n_configs, 2) - dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1692,42 +1667,80 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - op = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) - return qml.expval(op) + return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) res = _measure_operator() return res - assert np.allclose( - jax.jit(minimal_circ)(pars_q), jax.jit(jax.vmap(minimal_circ))(pars_q), tol - ) + res1 = jax.jit(minimal_circ)(pars_q) + res2 = jax.jit(jax.vmap(minimal_circ))(pars_q) + assert np.allclose(res1, res2, tol) def test_vmap_compared_param_broadcasting_multi_output( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled and when multiple output values are returned.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not yet support Hamiltonians") - elif diff_method == "hadamard": - pytest.skip("The Hadamard method does not yet support Hamiltonians") - if diff_method == "backprop": - pytest.skip( - "The backprop method does not yet support parameter-broadcasting with Hamiltonians" + n_configs = 5 + pars_q = np.random.rand(n_configs, 2) + + def minimal_circ(params): + @qml.qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + cache=None, ) + def _measure_operator(): + qml.RY(params[..., 0], wires=0) + qml.RY(params[..., 1], wires=1) + return qml.expval(qml.Z(0) @ qml.Z(1)), qml.expval(qml.X(0) @ qml.X(1)) + + res = _measure_operator() + return res + + res1, res2 = jax.jit(minimal_circ)(pars_q) + vres1, vres2 = jax.jit(jax.vmap(minimal_circ))(pars_q) + assert np.allclose(res1, vres1, tol) + assert np.allclose(res2, vres2, tol) + + def test_vmap_compared_param_broadcasting_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): + """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when + vectorized=True is specified for the callback when caching is disabled and when multiple output values + are returned.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") + elif dev.name == "lightning.qubit" and diff_method == "adjoint": + pytest.xfail("lightning adjoign cannot differentiate probabilities.") + interface = "jax-jit" - phys_qubits = 2 n_configs = 5 pars_q = np.random.rand(n_configs, 2) - dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1735,14 +1748,13 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - op1 = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) - op2 = qml.Hamiltonian([1.0], [qml.PauliX(0) @ qml.PauliX(1)]) - return qml.expval(op1), qml.expval(op2) + return qml.probs(wires=0), qml.probs(wires=1) res = _measure_operator() return res @@ -1758,23 +1770,25 @@ def _measure_operator(): @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestJIT: """Test JAX JIT integration with the QNode and automatic resolution of the correct JAX interface variant.""" - def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface): + def test_gradient( + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + ): """Test derivative calculation of a scalar valued QNode""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - gradient_kwargs = {} - if diff_method == "spsa": + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if device_vjp and jacobian == jax.jacfwd: + pytest.skip("device vjps not compatible with forward diff.") + elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1783,6 +1797,7 @@ def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x): @@ -1806,7 +1821,9 @@ def circuit(x): "ignore:Requested adjoint differentiation to be computed with finite shots." ) @pytest.mark.parametrize("shots", [10, 1000]) - def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface): + def test_hermitian( + self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface + ): """Test that the jax device works with qml.Hermitian and jitting even when shots>0. @@ -1814,13 +1831,6 @@ def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobi to different reasons, hence the parametrization in the test. """ # pylint: disable=unused-argument - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "backprop": pytest.skip("Backpropagation is unsupported if shots > 0.") @@ -1830,7 +1840,11 @@ def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobi projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1))) @qml.qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circ(projector): return qml.expval(qml.Hermitian(projector, wires=range(2))) @@ -1843,23 +1857,20 @@ def circ(projector): ) @pytest.mark.parametrize("shots", [10, 1000]) def test_probs_obs_none( - self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface ): """Test that the jax device works with qml.probs, a MeasurementProcess that has obs=None even when shots>0.""" # pylint: disable=unused-argument - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method in ["backprop", "adjoint"]: pytest.skip("Backpropagation is unsupported if shots > 0.") @qml.qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): return qml.probs(wires=0) @@ -1870,23 +1881,30 @@ def circuit(): # reason="Non-trainable parameters are not being correctly unwrapped by the interface" # ) def test_gradient_subset( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test derivative calculation of a scalar valued QNode with respect to a subset of arguments""" - if diff_method == "spsa" and not grad_on_execution: + if diff_method == "spsa" and not grad_on_execution and not device_vjp: pytest.xfail(reason="incorrect jacobian results") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if diff_method == "hadamard" and not grad_on_execution: - pytest.xfail(reason="XLA raised wire error") + if diff_method == "device" and not grad_on_execution and device_vjp: + pytest.xfail(reason="various runtime-related errors") + + if diff_method == "adjoint" and device_vjp and jacobian is jax.jacfwd: + pytest.xfail(reason="TypeError applying forward-mode autodiff.") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=1) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a, wires=0) @@ -1903,20 +1921,17 @@ def circuit(a, b, c): assert np.allclose(g, expected_g, atol=tol, rtol=0) def test_gradient_scalar_cost_vector_valued_qnode( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test derivative calculation of a scalar valued cost function that uses the output of a vector-valued QNode""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.xfail(reason="Adjoint does not support probs.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + elif jacobian == jax.jacfwd and device_vjp: + pytest.skip("device vjps are not compatible with forward differentiation.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1926,6 +1941,7 @@ def test_gradient_scalar_cost_vector_valued_qnode( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1958,21 +1974,26 @@ def cost(x, y, idx): assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0) assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0) + # pylint: disable=unused-argument def test_matrix_parameter( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test that the JAX-JIT interface works correctly with a matrix parameter""" - # pylint: disable=unused-argument - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("device vjps are not compatible with forward differentiation.") + # pylint: disable=unused-argument @qml.qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circ(p, U): qml.QubitUnitary(U, wires=0) @@ -1984,7 +2005,7 @@ def circ(p, U): res = jax.jit(circ)(p, U) assert np.allclose(res, -np.cos(p), atol=tol, rtol=0) - jac_fn = jax.jit(jax.grad(circ, argnums=(0))) + jac_fn = jax.jit(jacobian(circ, argnums=0)) res = jac_fn(p, U) assert np.allclose(res, np.sin(p), atol=tol, rtol=0) @@ -1992,27 +2013,31 @@ def circ(p, U): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturn: """Class to test the shape of the Grad/Jacobian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and one param, the gradient is a float.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2021,27 +2046,30 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.jit(jacobian(circuit))(a) + grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2051,7 +2079,9 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + grad = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + a, b, shots=shots + ) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -2059,21 +2089,24 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2082,31 +2115,31 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.jit(jacobian(circuit))(a) + grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2115,30 +2148,31 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2148,7 +2182,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) assert isinstance(jac, tuple) @@ -2159,24 +2193,25 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2184,29 +2219,33 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2214,7 +2253,9 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + par_0, par_1, shots=shots + ) assert isinstance(jac, tuple) @@ -2233,20 +2274,24 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2255,7 +2300,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2267,23 +2312,29 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - dev = qml.device(dev_name, wires=2, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if diff_method == "adjoint": + pytest.skip("adjoint supports either all expvals or only diagonal measurements") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2291,7 +2342,9 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + par_0, par_1, shots=shots + ) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2311,20 +2364,26 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - dev = qml.device(dev_name, wires=2, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if diff_method == "adjoint": + pytest.skip("adjoint supports either all expvals or only diagonal measurements") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2333,7 +2392,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2345,23 +2404,24 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if device_vjp and jacobian == jax.jacfwd: + pytest.skip("device vjp not compatible with forward differentiation.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2370,7 +2430,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2382,23 +2442,24 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2408,7 +2469,7 @@ def circuit(a, b): a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2428,23 +2489,24 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2453,7 +2515,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2474,23 +2536,17 @@ def circuit(a): @pytest.mark.parametrize("hessian", hessian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturnHessian: """Class to test the shape of the Hessian with different return types.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") par_0 = jax.numpy.array(0.1) @@ -2502,6 +2558,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2529,20 +2586,13 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2551,6 +2601,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2564,12 +2615,10 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2583,6 +2632,7 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2610,16 +2660,14 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") - dev = qml.device(dev_name, wires=2) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2628,6 +2676,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2641,17 +2690,10 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") @@ -2665,6 +2707,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2691,22 +2734,15 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_expval_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2715,6 +2751,7 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2733,12 +2770,10 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2752,6 +2787,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2778,16 +2814,14 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_var_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") - dev = qml.device(dev_name, wires=2) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2796,6 +2830,7 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2818,15 +2853,9 @@ def circuit(x): @pytest.mark.parametrize("diff_method", ["parameter-shift", "hadamard"]) def test_jax_device_hessian_shots(hessian, diff_method): """The hessian of multiple measurements with a multiple param array return a single array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device("default.qubit.jax", wires=num_wires, shots=10000) - - @jax.jit - @qml.qnode(dev, diff_method=diff_method, max_diff=2) + @partial(jax.jit, static_argnames="shots") + @qml.qnode(DefaultQubit(), diff_method=diff_method, max_diff=2) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -2835,7 +2864,7 @@ def circuit(x): x = jax.numpy.array([1.0, 2.0]) a, b = x - hess = jax.jit(hessian(circuit))(x) + hess = jax.jit(hessian(circuit), static_argnames="shots")(x, shots=10000) expected_hess = [ [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], @@ -2849,28 +2878,33 @@ def circuit(x): @pytest.mark.parametrize("argnums", [0, 1, [0, 1]]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestSubsetArgnums: def test_single_measurement( self, interface, - dev_name, + dev, diff_method, grad_on_execution, + device_vjp, jacobian, argnums, jit_inside, tol, ): """Test single measurement with different diff methods with argnums.""" - - dev = qml.device(dev_name, wires=3) - kwargs = {} + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transform have a different vjp shape convention.") if diff_method == "spsa": - tol = TOL_FOR_SPSA kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA @qml.qnode( dev, @@ -2878,6 +2912,7 @@ def test_single_measurement( diff_method=diff_method, grad_on_execution=grad_on_execution, cache=False, + device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2907,27 +2942,34 @@ def circuit(a, b): def test_multi_measurements( self, interface, - dev_name, + dev, diff_method, grad_on_execution, + device_vjp, jacobian, argnums, jit_inside, tol, ): """Test multiple measurements with different diff methods with argnums.""" - dev = qml.device(dev_name, wires=3) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transform have a different vjp shape convention.") kwargs = {} if diff_method == "spsa": - tol = TOL_FOR_SPSA kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA @qml.qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2953,3 +2995,72 @@ def circuit(a, b): else: assert np.allclose(jac[0], expected[0], atol=tol) assert np.allclose(jac[1], expected[1], atol=tol) + + +class TestSinglePrecision: + """Tests for compatibility with single precision mode.""" + + # pylint: disable=import-outside-toplevel + def test_type_conversion_fallback(self): + """Test that if the type isn't int, float, or complex, we still have a fallback.""" + from pennylane.workflow.interfaces.jax_jit import _jax_dtype + + assert _jax_dtype(bool) == jax.numpy.dtype(bool) + + @pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) + def test_float32_return(self, diff_method): + """Test that jax jit works when float64 mode is disabled.""" + jax.config.update("jax_enable_x64", False) + + try: + + @jax.jit + @qml.qnode(qml.device("default.qubit"), diff_method=diff_method) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + grad = jax.grad(circuit)(jax.numpy.array(0.1)) + assert qml.math.allclose(grad, -np.sin(0.1)) + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) + + @pytest.mark.parametrize("diff_method", ("adjoint", "finite-diff")) + def test_complex64_return(self, diff_method): + """Test that jax jit works with differentiating the state.""" + jax.config.update("jax_enable_x64", False) + + try: + tol = 2e-2 if diff_method == "finite-diff" else 1e-6 + + @jax.jit + @qml.qnode(qml.device("default.qubit", wires=1), diff_method=diff_method) + def circuit(x): + qml.RX(x, wires=0) + return qml.state() + + j = jax.jacobian(circuit, holomorphic=True)(jax.numpy.array(0.1 + 0j)) + assert qml.math.allclose(j, [-np.sin(0.05) / 2, -np.cos(0.05) / 2 * 1j], atol=tol) + + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) + + def test_int32_return(self): + """Test that jax jit forward execution works with samples and int32""" + + jax.config.update("jax_enable_x64", False) + + try: + + @jax.jit + @qml.qnode(qml.device("default.qubit", shots=10), diff_method=qml.gradients.param_shift) + def circuit(x): + qml.RX(x, wires=0) + return qml.sample(wires=0) + + _ = circuit(jax.numpy.array(0.1)) + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py index fa88e315155..3ea650f96e8 100644 --- a/tests/interfaces/test_jax_qnode.py +++ b/tests/interfaces/test_jax_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,27 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-Python interface with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods,too-many-public-methods +# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-implicit-booleaness-not-comparison + +from itertools import product + +import numpy as np import pytest import pennylane as qml -from pennylane import numpy as np from pennylane import qnode -from pennylane.tape import QuantumScript - -qubit_device_and_diff_method = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], +from pennylane.devices import DefaultQubit + +device_seed = 42 + +# device, diff_method, grad_on_execution, device_vjp +device_and_diff_method = [ + [DefaultQubit(seed=device_seed), "backprop", True, False], + [DefaultQubit(seed=device_seed), "finite-diff", False, False], + [DefaultQubit(seed=device_seed), "parameter-shift", False, False], + [DefaultQubit(seed=device_seed), "adjoint", True, False], + [DefaultQubit(seed=device_seed), "adjoint", False, False], + [DefaultQubit(seed=device_seed), "adjoint", True, True], + [DefaultQubit(seed=device_seed), "adjoint", False, True], + [DefaultQubit(seed=device_seed), "spsa", False, False], + [DefaultQubit(seed=device_seed), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], ] -interface_and_qubit_device_and_diff_method = [ - ["auto"] + inner_list for inner_list in qubit_device_and_diff_method -] + [["jax"] + inner_list for inner_list in qubit_device_and_diff_method] +interface_and_device_and_diff_method = [ + ["auto"] + inner_list for inner_list in device_and_diff_method +] + [["jax"] + inner_list for inner_list in device_and_diff_method] pytestmark = pytest.mark.jax @@ -46,46 +58,49 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", interface_and_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, interface, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = np.array(0.1, requires_grad=True) + a = jax.numpy.array(0.1) circuit(a) assert circuit.interface == interface - # the tape is able to deduce trainable parameters - assert circuit.qtape.trainable_params == [0] + # jax doesn't set trainable parameters on regular execution + assert circuit.qtape.trainable_params == [] # gradients should work grad = jax.grad(circuit)(a) assert isinstance(grad, jax.Array) + # the tape is able to deduce trainable parameters + assert circuit.qtape.trainable_params == [0] assert grad.shape == () - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): # pylint:disable=unused-argument """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -94,14 +109,7 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, i a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=2) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface=interface, diff_method="parameter-shift") def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -127,27 +135,18 @@ def circuit(a, b): expected = [-np.sin(a) + np.sin(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - # trainability also updates on evaluation - a = np.array(0.54, requires_grad=False) - b = np.array(0.8, requires_grad=True) - circuit(a, b) - assert circuit.qtape.trainable_params == [1] - - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -162,21 +161,20 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test that the jax interface works correctly with a matrix parameter""" U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -189,39 +187,32 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test that operation and nested tape expansion is differentiable""" + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) if diff_method == "spsa": - spsa_kwargs = dict( - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), - num_directions=10, - ) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA - class U3(qml.U3): + class U3(qml.U3): # pylint:disable=too-few-public-methods def decomposition(self): theta, phi, lam = self.data wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] - with qml.queuing.AnnotatedQueue() as q_tape: - qml.Rot(lam, theta, -lam, wires=wires) - qml.PhaseShift(phi + lam, wires=wires) - - tape = QuantumScript.from_queue(q_tape) - return tape.operations - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -250,23 +241,16 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): # pylint:disable=unused-argument """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device(dev_name, wires=1) + a = jax.numpy.array([0.1, 0.2]) - @qnode( - dev, - interface=interface, - diff_method="finite-diff", - h=1e-8, - approx_order=2, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface=interface, diff_method="finite-diff", h=1e-8, approx_order=2) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -276,30 +260,31 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_expval( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) @qnode(dev, **kwargs) def circuit(a, b): @@ -310,7 +295,7 @@ def circuit(a, b): res = circuit(a, b) - assert circuit.qtape.trainable_params == [0, 1] + assert circuit.qtape.trainable_params == [] assert isinstance(res, tuple) assert len(res) == 2 @@ -320,6 +305,7 @@ def circuit(a, b): res = jax.jacobian(circuit, argnums=[0, 1])(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) + assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -339,11 +325,20 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_jacobian_no_evaluate( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA @@ -351,13 +346,6 @@ def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, in a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -373,11 +361,10 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for _res, _exp in zip(res, expected): - for r, e in zip(_res, _exp): - assert isinstance(r, jax.numpy.ndarray) - assert r.shape == () - assert np.allclose(r, e, atol=tol, rtol=0) + assert isinstance(res[0][0], jax.numpy.ndarray) + for i, j in product((0, 1), (0, 1)): + assert res[i][j].shape == () + assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) # call the Jacobian with new parameters a = jax.numpy.array(0.6) @@ -390,30 +377,27 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for _res, _exp in zip(res, expected): - for r, e in zip(_res, _exp): - assert isinstance(r, jax.numpy.ndarray) - assert r.shape == () - assert np.allclose(r, e, atol=tol, rtol=0) + for i, j in product((0, 1), (0, 1)): + assert res[i][j].shape == () + assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) - def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_single_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with a single prob output""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -445,24 +429,24 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_multi_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -484,11 +468,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) + assert res[0].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (4,) + assert res[1].shape == (4,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -527,24 +511,23 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -561,11 +544,11 @@ def circuit(x, y): assert len(res) == 2 assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () + assert res[0].shape == () # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) + assert res[1].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -599,32 +582,28 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - tol = TOL_FOR_SPSA + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + tol = TOL_FOR_SPSA x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) + if diff_method == "adjoint": + x = x + 0j + y = y + 0j + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -633,6 +612,8 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) + if "lightning" in dev.name: + pytest.xfail("lightning does not support measuring probabilities with adjoint.") jac = jax.jacobian(circuit, argnums=[0])(x, y) expected = [ @@ -657,21 +638,24 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "hadamard": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if diff_method == "hadamard": pytest.skip("Hadamard does not support var") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -690,11 +674,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () + assert res[0].shape == () # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) + assert res[1].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -734,53 +718,40 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test jax device works with diff_method=None.""" - dev = qml.device("default.qubit.jax", wires=1, shots=10) + """Test device works with diff_method=None.""" - @qml.qnode(dev, diff_method=None, interface=interface) + @qml.qnode(DefaultQubit(), diff_method=None, interface=interface) def circuit(x): qml.RX(x, wires=0) return qml.expval(qml.PauliZ(0)) - assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) + assert jax.numpy.allclose(circuit(jax.numpy.array(0.0), shots=10), 1) - def test_changing_shots(self, interface, mocker, tol): + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # no additional calls + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # pylint: disable=comparison-with-callable def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -788,51 +759,40 @@ def cost_fn(a, b): return qml.expval(qml.PauliY(1)) res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) - assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) - def test_update_diff_method(self, mocker, interface): + def test_update_diff_method(self, interface, mocker): """Test that temporarily setting the shots updates the diff method""" - # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - # We're choosing interface="jax" such that backprop can be used in the - # test later - @qnode(dev, interface="jax") + @qnode(DefaultQubit(), interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - assert cost_fn.gradient_fn is qml.gradients.param_shift - - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method,grad_on_execution, device_vjp", device_and_diff_method) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev_name, diff_method, grad_on_execution): + def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): """Test sampling works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") @@ -840,43 +800,48 @@ def test_sampling(self, dev_name, diff_method, grad_on_execution): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) + @qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) + assert res[0].shape == (10,) # pylint:disable=comparison-with-callable assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) + assert res[1].shape == (10,) # pylint:disable=comparison-with-callable - def test_counts(self, dev_name, diff_method, grad_on_execution): + def test_counts(self, dev, diff_method, grad_on_execution, device_vjp): """Test counts works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") if diff_method == "adjoint": - pytest.skip("Adjoint warns with finite shots") + pytest.skip("Adjoint errors with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) + @qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return ( - qml.counts(qml.PauliZ(0), all_outcomes=True), - qml.counts(qml.PauliX(1), all_outcomes=True), - ) + return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) @@ -885,25 +850,32 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + # pylint:disable=too-few-public-methods class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="jax", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) - @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="jax", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -928,37 +900,72 @@ def cost(weights): assert len(res) == 2 + def test_postselection_differentiation(self, dev, diff_method, grad_on_execution, device_vjp): + """Test that when postselecting with default.qubit, differentiation works correctly.""" + + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") + + @qml.qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + phi = jax.numpy.array(1.23) + theta = jax.numpy.array(4.56) + + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + + gradient = jax.grad(circuit, argnums=[0, 1])(phi, theta) + exp_theta_grad = jax.grad(expected_circuit)(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) + @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test second derivative calculation of a scalar-valued QNode""" - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - max_diff=2, - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + "max_diff": 2, + } + if diff_method == "adjoint": pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - spsa_kwargs = dict( - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=100, h=0.001 - ) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(x): qml.RY(x[0], wires=0) @@ -987,7 +994,7 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": @@ -1000,18 +1007,12 @@ def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol) } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1045,12 +1046,15 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": + qml.math.random.seed(42) gradient_kwargs = { "h": H_FOR_SPSA, "num_directions": 20, @@ -1058,18 +1062,12 @@ def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, i } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1113,7 +1111,7 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, interface, grad_on_execution, tol + self, dev, diff_method, interface, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} @@ -1127,18 +1125,12 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1185,7 +1177,7 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} @@ -1199,18 +1191,12 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1236,6 +1222,7 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) + hess = jax.jacobian(jac_fn, argnums=[0, 1])(a, b) expected_hess = np.array( @@ -1255,23 +1242,21 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning does not support state adjoint differentiation.") x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1281,7 +1266,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") + assert res.dtype is np.dtype("complex128") # pylint:disable=no-member probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1295,18 +1280,19 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "hadamard": + pytest.skip("adjoint supports all expvals or only diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Hadamard does not support var.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1315,6 +1301,7 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1338,94 +1325,7 @@ def circuit(x, y): @pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -@pytest.mark.parametrize("interface", ["jax", "jax-python"]) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = 0.543 - phi = -0.654 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(r, phi) - expected = np.array( - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - n = 0.12 - a = 0.765 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(n, a) - expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"]) -def test_adjoint_reuse_device_state(mocker, interface): - """Tests that the jax interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - jax.grad(circ)(1.0) - assert circ.device.num_executions == 1 - - spy.assert_called_with(mocker.ANY, use_device_state=True) - - -@pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1433,21 +1333,14 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1459,6 +1352,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1474,7 +1368,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol + self, dev, diff_method, grad_on_execution, max_diff, interface, device_vjp, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None @@ -1492,7 +1386,6 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1501,6 +1394,7 @@ def test_hamiltonian_expansion_analytic( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1545,11 +1439,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker ): - """Test that the Hamiltonian is expanded if there - are non-commuting groups and the number of shots is finite + """Test that the Hamiltonian is correctly measured (and not expanded) + if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 @@ -1565,7 +1459,6 @@ def test_hamiltonian_expansion_finite_shots( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=50000) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1574,6 +1467,7 @@ def test_hamiltonian_expansion_finite_shots( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1590,13 +1484,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() + spy.assert_not_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1626,27 +1520,24 @@ def circuit(data, weights, coeffs): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) -class TestReturn: +class TestReturn: # pylint:disable=too-many-public-methods """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1655,27 +1546,24 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.grad(circuit)(a) + grad = jax.grad(circuit)(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, shots, device_vjp, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1685,7 +1573,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.grad(circuit, argnums=[0, 1])(a, b) + grad = jax.grad(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -1693,21 +1581,18 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, shots, device_vjp, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1716,14 +1601,14 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.grad(circuit)(a) + grad = jax.grad(circuit)(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1733,15 +1618,12 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1750,14 +1632,14 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1766,15 +1648,12 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1784,7 +1663,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b) + jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(jac, tuple) @@ -1796,7 +1675,7 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1805,15 +1684,12 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1821,26 +1697,29 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, + dev, + diff_method, + grad_on_execution, + jacobian, + shots, + interface, + device_vjp, ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - + if device_vjp and jacobian is jax.jacfwd: + pytest.skip("forward pass can't be done with registered vjp.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1850,6 +1729,7 @@ def test_jacobian_expval_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1857,7 +1737,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(jac, tuple) @@ -1877,21 +1757,22 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if device_vjp and jacobian is jax.jacfwd: + pytest.skip("forward pass can't be done with registered vjp.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1900,7 +1781,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1913,18 +1794,16 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all measurements or only diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1934,6 +1813,7 @@ def test_jacobian_var_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1941,7 +1821,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1962,20 +1842,22 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all expvals or all diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1984,7 +1866,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1997,24 +1879,22 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2023,7 +1903,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2036,33 +1916,32 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b) + jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2083,23 +1962,22 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2108,7 +1986,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2120,19 +1998,12 @@ def circuit(a): assert jac[1].shape == (4, 2) def test_hessian_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2145,6 +2016,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2152,7 +2024,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2170,7 +2042,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2178,13 +2050,6 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2193,6 +2058,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2200,13 +2066,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2215,7 +2081,6 @@ def test_hessian_var_multiple_params( pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -2226,6 +2091,7 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2233,7 +2099,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2251,7 +2117,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2261,8 +2127,6 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2271,6 +2135,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2278,22 +2143,15 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2311,6 +2169,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2318,7 +2177,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2354,7 +2213,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2364,19 +2223,13 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( dev, interface=interface, diff_method=diff_method, + device_vjp=device_vjp, max_diff=2, grad_on_execution=grad_on_execution, ) @@ -2386,7 +2239,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2398,7 +2251,7 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2408,8 +2261,6 @@ def test_hessian_probs_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = qml.numpy.array(0.1) par_1 = qml.numpy.array(0.2) @@ -2419,6 +2270,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2426,7 +2278,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2462,7 +2314,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_var_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2472,8 +2324,6 @@ def test_hessian_var_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2482,6 +2332,7 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2489,7 +2340,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2501,14 +2352,11 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="jax") + @qml.qnode(DefaultQubit(), interface="jax") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_jax_qnode_shot_vector.py b/tests/interfaces/test_jax_qnode_shot_vector.py index 0ef5c8291de..697a1e90223 100644 --- a/tests/interfaces/test_jax_qnode_shot_vector.py +++ b/tests/interfaces/test_jax_qnode_shot_vector.py @@ -28,9 +28,9 @@ all_shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)] qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], + ["default.qubit", "finite-diff", {"h": 10e-2}], + ["default.qubit", "parameter-shift", {}], + ["default.qubit", "spsa", {"h": 10e-2, "num_directions": 20}], ] interface_and_qubit_device_and_diff_method = [ @@ -766,59 +766,9 @@ def circuit(x): assert h[1].shape == (2, 2, 2) -@pytest.mark.parametrize("shots", all_shots) -class TestReturnShotVectorsDevice: - """Test for shot vectors with device method adjoint_jacobian.""" - - def test_jac_adjoint_fwd_error(self, shots): - """Test that an error is raised for adjoint forward.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - with pytest.warns( - UserWarning, match="Requested adjoint differentiation to be computed with finite shots." - ): - - @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=True) - def circuit(a): - qml.RY(a, wires=0) - qml.RX(0.2, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = jax.numpy.array(0.1) - - if isinstance(shots, tuple): - with pytest.raises( - qml.QuantumFunctionError, - match="Adjoint does not support shot vectors.", - ): - jax.jacobian(circuit)(a) - - def test_jac_adjoint_bwd_error(self, shots): - """Test that an error is raised for adjoint backward.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - with pytest.warns( - UserWarning, match="Requested adjoint differentiation to be computed with finite shots." - ): - - @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=False) - def circuit(a): - qml.RY(a, wires=0) - qml.RX(0.2, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = jax.numpy.array(0.1) - - with pytest.raises( - qml.QuantumFunctionError, - match="Adjoint does not support shot vectors.", - ): - jax.jacobian(circuit)(a) - - qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit", "finite-diff", {"h": 10e-2}], + ["default.qubit", "parameter-shift", {}], ] shots_large = [(1000000, 900000, 800000), (1000000, (900000, 2))] diff --git a/tests/interfaces/test_set_shots.py b/tests/interfaces/test_set_shots.py index 5e6a53b83d3..efe8c9cf7ca 100644 --- a/tests/interfaces/test_set_shots.py +++ b/tests/interfaces/test_set_shots.py @@ -18,7 +18,6 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots from pennylane.workflow import set_shots @@ -30,29 +29,3 @@ def test_shots_new_device_interface(): with pytest.raises(ValueError): with set_shots(dev, 10): pass - - -def test_set_with_shots_class(): - """Test that shots can be set on the old device interface with a Shots class.""" - - dev = qml.devices.DefaultQubitLegacy(wires=1) - with set_shots(dev, Shots(10)): - assert dev.shots == 10 - - assert dev.shots is None - - shot_tuples = Shots((10, 10)) - with set_shots(dev, shot_tuples): - assert dev.shots == 20 - assert dev.shot_vector == list(shot_tuples.shot_vector) - - assert dev.shots is None - - -def test_shots_not_altered_if_False(): - """Test a value of False can be passed to shots, indicating to not override - shots on the device.""" - - dev = qml.devices.DefaultQubitLegacy(wires=1) - with set_shots(dev, False): - assert dev.shots is None diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py index 057f93d0038..11a54575cfa 100644 --- a/tests/interfaces/test_tensorflow.py +++ b/tests/interfaces/test_tensorflow.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,227 +11,32 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the TensorFlow interface""" -# pylint: disable=protected-access,too-few-public-methods +"""Tensorflow specific tests for execute and default qubit 2.""" import numpy as np import pytest import pennylane as qml from pennylane import execute -from pennylane.gradients import finite_diff, param_shift +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift pytestmark = pytest.mark.tf - -tf = pytest.importorskip("tensorflow", minversion="2.1") - - -class TestTensorFlowExecuteUnitTests: - """Unit tests for TensorFlow execution""" - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = tf.Variable([0.1, 0.2], dtype=tf.float64) - - dev = qml.device("default.qubit.legacy", wires=1) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - interface="tf", - )[0] - - res = t.jacobian(res, a) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if a gradient transform - is used with grad_on_execution""" - a = tf.Variable([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - execute([tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface="tf") - - def test_grad_on_execution(self, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable([0.1, 0.2]) - spy = mocker.spy(dev, "execute_and_gradients") - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface="tf", - ) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_no_grad_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - a = tf.Variable([0.1, 0.2]) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface="tf", - )[0] - - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - - t.jacobian(res, a) - spy_gradients.assert_called() +tf = pytest.importorskip("tensorflow") +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - a = tf.Variable([0.1, 0.2]) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, gradient_fn=param_shift, cachesize=2, interface="tf")[0] - - t.jacobian(res, a) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - a = tf.Variable([0.1, 0.2]) - custom_cache = {} - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, gradient_fn=param_shift, cache=custom_cache, interface="tf")[ - 0 - ] - - t.jacobian(res, a) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - unwrapped_tape = qml.transforms.convert_to_numpy_parameters(tape)[0][0] - h = unwrapped_tape.hash - - assert h in cache - assert np.allclose(cache[h], res) - - def test_caching_param_shift(self): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum.""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="tf")[0] - - # Without caching, and non-vectorized, 9 evaluations are required to compute - # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) - with tf.GradientTape(persistent=True) as t: - res = cost(a, cache=None) - t.jacobian(res, a, experimental_use_pfor=False) - assert dev.num_executions == 9 - - # With caching, and non-vectorized, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - with tf.GradientTape(persistent=True) as t: - res = cost(a, cache=True) - t.jacobian(res, a) - assert dev.num_executions == 5 - - # In vectorized mode, 5 evaluations are required to compute - # the Jacobian regardless of caching: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - with tf.GradientTape() as t: - res = cost(a, cache=None) - t.jacobian(res, a) - assert dev.num_executions == 5 + """Tests for caching behaviour""" @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): + def test_caching_param_shift_hessian(self, num_params): """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = tf.Variable(np.arange(1, num_params + 1) / 10, dtype=tf.float64) + dev = DefaultQubit() + params = tf.Variable(tf.range(1, num_params + 1) / 10) - N = params.shape[0] + N = num_params def cost(x, cache): with qml.queuing.AnnotatedQueue() as q: @@ -242,362 +47,425 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], dev, gradient_fn=param_shift, cache=cache, interface="tf", max_diff=2 + return qml.execute( + [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 )[0] # No caching: number of executions is not ideal - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - res = cost(params, cache=False) - grad = t1.gradient(res, params) - hess1 = t2.jacobian(grad, params) + with qml.Tracker(dev) as tracker: + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost(params, cache=False) + grad = grad_tape.gradient(res, params) + hess1 = jac_tape.jacobian(grad, params) if num_params == 2: # compare to theoretical result - x, y, *_ = params * 1.0 - expected = np.array( + x, y, *_ = params + expected = tf.convert_to_tensor( [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [2 * tf.cos(2 * x) * tf.sin(y) ** 2, tf.sin(2 * x) * tf.sin(2 * y)], + [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert np.allclose(expected, hess1) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian - nonideal_runs = dev.num_executions + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - res = cost(params, cache=True) - grad = t1.gradient(res, params) - hess2 = t2.jacobian(grad, params) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + with qml.Tracker(dev) as tracker2: + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost(params, cache=True) + grad = grad_tape.gradient(res, params) + hess2 = jac_tape.jacobian(grad, params) + assert np.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal - assert expected_runs_ideal < nonideal_runs - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift, "interface": "tf"}, - {"gradient_fn": param_shift, "interface": "auto"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "tf", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "tf", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "auto", - }, + assert tracker2.totals["executions"] == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)), # 0 + ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()), # 1 + ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()), # 2 + ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()), # 3 + ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)), # 4 + ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()), # 5 + ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()), # 6 + ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()), # 7 + ({"gradient_fn": "adjoint", "interface": "tf", "device_vjp": True}, None, DefaultQubit()), # 8 ] -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) -class TestTensorFlowExecuteIntegration: - """Test the TensorFlow interface execute function +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 + + +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +class TestTensorflowExecuteIntegration: + """Test the tensorflow interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable(0.1) - b = tf.Variable(0.2) - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b): + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) + + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + return execute([tape1, tape2], device, **execute_kwargs) - tape2 = qml.tape.QuantumScript.from_queue(q2) - res = execute([tape1, tape2], dev, **execute_kwargs) + a = tf.Variable(0.1, dtype="float64") + b = tf.constant(0.2, dtype="float64") + with device.tracker: + res = cost(a, b) + + if execute_kwargs.get("gradient_fn", None) == "adjoint" and not execute_kwargs.get( + "device_vjp", False + ): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 assert res[0].shape == () assert res[1].shape == () - assert isinstance(res[0], tf.Tensor) - assert isinstance(res[1], tf.Tensor) - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" a = tf.Variable(0.1, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] + device_vjp = execute_kwargs.get("device_vjp", False) - res = t.jacobian(res, a) - assert res.shape == () + def cost(a): + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(a) + res = tape.jacobian(cost_res, a, experimental_use_pfor=not device_vjp) + assert res.shape == () # pylint: disable=no-member - tape = qml.tape.QuantumScript.from_queue(q) + # compare to standard tape jacobian + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -tf.sin(a), atol=atol_for_shots(shots)) - def test_jacobian(self, execute_kwargs, tol): + def test_jacobian(self, execute_kwargs, shots, device): """Test jacobian calculation""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, max_diff=2, **execute_kwargs)[0] - res = tf.stack(res) + device_vjp = execute_kwargs.get("device_vjp", False) - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(a, b) + expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - (agrad, bgrad) = t.jacobian(res, [a, b]) - assert agrad.shape == (2,) - assert bgrad.shape == (2,) + jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + assert isinstance(jac, list) and len(jac) == 2 + assert jac[0].shape == (2,) + assert jac[1].shape == (2,) - expected = [[-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]] - assert np.allclose(expected, [agrad, bgrad], atol=tol, rtol=0) + expected = ([-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]) + for _r, _e in zip(jac, expected): + assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - def test_tape_no_parameters(self, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = 1.0 * params - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(0.5, wires=0) - qml.expval(qml.PauliZ(0)) + tape2 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape3 = qml.tape.QuantumScript.from_queue(q3) - res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) - res = tf.stack(res) + tape4 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, + ) + return tf.reduce_sum( + qml.math.hstack( + execute([tape1, tape2, tape3, tape4], device, **execute_kwargs), + like="tensorflow", + ) + ) - expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert np.allclose(res, expected, atol=tol, rtol=0) + params = tf.Variable([0.1, 0.2]) + x, y = params - grad = t.gradient(res, params) - expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + with tf.GradientTape() as tape: + res = cost(params) + expected = 2 + tf.cos(0.5) + tf.cos(x) * tf.cos(y) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_reusing_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + if ( + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" + ): + with pytest.raises(NotImplementedError): + tape.gradient(res, params) + return - dev = qml.device("default.qubit.legacy", wires=2) + grad = tape.gradient(res, params) + expected = [-tf.cos(y) * tf.sin(x), -tf.cos(x) * tf.sin(y)] + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - tape = qml.tape.QuantumScript.from_queue(q) - assert tape.trainable_params == [0, 1] - res = execute([tape], dev, **execute_kwargs)[0] - res = tf.stack(res) + if ( + execute_kwargs["gradient_fn"] == "adjoint" + and execute_kwargs["interface"] == "tf-autograph" + ): + pytest.skip("Cannot compute the jacobian with adjoint-differentation and tf-autograph") - t.jacobian(res, [a, b]) + device_vjp = execute_kwargs.get("device_vjp", False) - a = tf.Variable(0.54, dtype=tf.float64) - b = tf.Variable(0.8, dtype=tf.float64) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] - res2 = tf.stack(res2) + tape2 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5, dtype=tf.float64), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + return qml.math.hstack( + execute([tape1, tape2, tape3], device, **execute_kwargs), like="tensorflow" + ) - jac2 = t.jacobian(res2, [a, b]) - expected = [ - [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], - [0, -tf.cos(2 * a) * tf.cos(b)], - ] - assert np.allclose(jac2, expected, atol=tol, rtol=0) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = params - def test_reusing_pre_constructed_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape that was previously constructed - *outside of* a gradient tape, by passing new parameters""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(params) + + assert isinstance(res, tf.Tensor) + assert res.shape == (4,) + + assert np.allclose(res[0], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], tf.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) + + jac = tape.jacobian(res, params, experimental_use_pfor=not device_vjp) + assert isinstance(jac, tf.Tensor) + assert jac.shape == (4, 2) # pylint: disable=no-member + + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - dev = qml.device("default.qubit.legacy", wires=2) + d1 = -tf.sin(x) * tf.cos(y) + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) + + d2 = -tf.cos(x) * tf.sin(y) + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + a = tf.Variable(0.1) + b = tf.Variable(0.2) + + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + ) + assert tape.trainable_params == [0, 1] - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) + device_vjp = execute_kwargs.get("device_vjp", False) - tape = qml.tape.QuantumScript.from_queue(q) - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([a, b], [0, 1]) - assert tape.trainable_params == [0, 1] - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return qml.math.hstack( + execute([new_tape], device, **execute_kwargs)[0], like="tensorflow" + ) - t.jacobian(res, [a, b]) + with tf.GradientTape(persistent=device_vjp) as t: + res = cost(a, b) + jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) a = tf.Variable(0.54, dtype=tf.float64) b = tf.Variable(0.8, dtype=tf.float64) - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] - res2 = qml.math.stack(res2) + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + + with tf.GradientTape(persistent=device_vjp): + res2 = cost(2 * a, b) expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - jac2 = t.jacobian(res2, [a, b]) - expected = [ + with tf.GradientTape(persistent=device_vjp) as t: + res = cost(2 * a, b) + + jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + expected = ( [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], [0, -tf.cos(2 * a) * tf.cos(b)], - ] - assert np.allclose(jac2, expected, atol=tol, rtol=0) + ) + assert isinstance(jac, list) and len(jac) == 2 + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing(self, execute_kwargs): + def test_classical_processing(self, execute_kwargs, device, shots): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=1) + device_vjp = execute_kwargs.get("device_vjp", False) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + tf.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + tf.sin(a), wires=0), + ] - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [0, 2] - assert tape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] + + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(a, b, c) + + res = tape.jacobian(cost_res, [a, c], experimental_use_pfor=not device_vjp) + + # Only two arguments are trainable + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () - res = t.jacobian(res, [a, b, c]) - assert isinstance(res[0], tf.Tensor) - assert res[1] is None - assert isinstance(res[2], tf.Tensor) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - def test_no_trainable_parameters(self, execute_kwargs): + def test_no_trainable_parameters(self, execute_kwargs, shots, device): """Test evaluation and Jacobian if there are no trainable parameters""" - b = tf.constant(0.2, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + a = tf.constant(0.1) + b = tf.constant(0.2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(0.2, wires=0) - qml.RX(b, wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + with tf.GradientTape() as tape: + cost_res = cost(a, b) - assert res.shape == (2,) - assert isinstance(res, tf.Tensor) + assert cost_res.shape == (2,) + + res = tape.jacobian(cost_res, [a, b]) + assert len(res) == 2 + assert all(r is None for r in res) + + def loss(a, b): + return tf.reduce_sum(cost(a, b)) + + with tf.GradientTape() as tape: + loss_res = loss(a, b) - res = t.jacobian(res, b) - assert res is None + res = tape.gradient(loss_res, [a, b]) + assert all(r is None for r in res) - @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter(self, execute_kwargs, U, tol): - """Test that the TF interface works correctly + def test_matrix_parameter(self, execute_kwargs, device, shots): + """Test that the tensorflow interface works correctly with a matrix parameter""" - a = tf.Variable(0.1, dtype=tf.float64) + U = tf.constant([[0, 1], [1, 0]], dtype=tf.complex128) + a = tf.Variable(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + device_vjp = execute_kwargs.get("device_vjp", False) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1] + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(a, U) - assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) + assert np.allclose(res, -tf.cos(a), atol=atol_for_shots(shots)) - res = t.jacobian(res, a) - assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + assert isinstance(jac, tf.Tensor) + assert np.allclose(jac, tf.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): - """Test that operation and nested tape expansion + def test_differentiable_expand(self, execute_kwargs, device, shots): + """Test that operation and nested tapes expansion is differentiable""" + device_vjp = execute_kwargs.get("device_vjp", False) + class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -606,26 +474,38 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device("default.qubit.legacy", wires=1) - a = np.array(0.1) - p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - with tf.GradientTape() as tape: - with qml.queuing.AnnotatedQueue() as q_qtape: - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - qml.expval(qml.PauliX(0)) + a = tf.constant(0.1) + p = tf.Variable([0.1, 0.2, 0.3]) - qtape = qml.tape.QuantumScript.from_queue(q_qtape) - res = execute([qtape], dev, **execute_kwargs)[0] + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost_fn(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = tape.jacobian(res, p) - expected = np.array( + res = tape.jacobian(cost_res, p, experimental_use_pfor=not device_vjp) + expected = tf.convert_to_tensor( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) @@ -635,121 +515,119 @@ def decomposition(self): * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + + device_vjp = execute_kwargs.get("device_vjp", False) - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(x, y) - expected = np.array( + expected = tf.convert_to_tensor( [ - [tf.cos(x / 2) ** 2, tf.sin(x / 2) ** 2], - [(1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2], + [ + tf.cos(x / 2) ** 2, + tf.sin(x / 2) ** 2, + (1 + tf.cos(x) * tf.cos(y)) / 2, + (1 - tf.cos(x) * tf.cos(y)) / 2, + ], ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = t.jacobian(res, [x, y]) - expected = np.array( - [ + if ( + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" + ): + with pytest.raises(tf.errors.UnimplementedError): + tape.jacobian(cost_res, [x, y]) + return + res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ( + tf.convert_to_tensor( [ - [-tf.sin(x) / 2, tf.sin(x) / 2], - [-tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], - ], + [ + -tf.sin(x) / 2, + tf.sin(x) / 2, + -tf.sin(x) * tf.cos(y) / 2, + tf.sin(x) * tf.cos(y) / 2, + ], + ] + ), + tf.convert_to_tensor( [ - [0, 0], - [-tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ], - ] + [0, 0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ] + ), ) - assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation(self, execute_kwargs, tol): + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - dev = qml.device("default.qubit.legacy", wires=2) + device_vjp = execute_kwargs.get("device_vjp", False) + + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = tf.experimental.numpy.hstack(res) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(x, y) - expected = np.array( + expected = tf.convert_to_tensor( [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = t.jacobian(res, [x, y]) - expected = np.array( - [ - [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], - [0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_sampling(self, execute_kwargs): - """Test sampling works as expected""" if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" ): - pytest.skip("Adjoint differentiation does not support samples") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(tf.Variable(0.1), wires=0) - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) - - assert res.shape == (2, 10) - assert isinstance(res, tf.Tensor) + with pytest.raises(tf.errors.UnimplementedError): + tape.jacobian(cost_res, [x, y]) + return + res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = ( + tf.convert_to_tensor( + [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.sin(x) * tf.cos(y) / 2] + ), + tf.convert_to_tensor([0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2]), + ) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) -@pytest.mark.parametrize("interface", ["auto", "tf"]) class TestHigherOrderDerivatives: - """Test that the TensorFlow execute function can be differentiated""" + """Test that the tensorflow execute function can be differentiated""" - @pytest.mark.slow @pytest.mark.parametrize( "params", [ @@ -758,196 +636,96 @@ class TestHigherOrderDerivatives: tf.Variable([-2.0, 0], dtype=tf.float64), ], ) - def test_parameter_shift_hessian(self, params, tol, interface): + def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using tensorflow, yielding second derivatives.""" - dev = qml.device("default.qubit.tf", wires=2) - x, y = params * 1.0 - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(params[0], wires=0) - qml.RY(params[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, interface=interface, max_diff=2 - ) - res = result[0] + result[1][0] + dev = DefaultQubit() - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost_fn(x): + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - grad = t1.gradient(res, params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + return result[0] + result[1][0] - hess = t2.jacobian(grad, params) - expected = np.array( - [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(hess, expected, atol=tol, rtol=0) - - def test_hessian_vector_valued(self, tol, interface): - """Test hessian calculation of a vector valued QNode""" - dev = qml.device("default.qubit.tf", wires=1) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - - with tf.GradientTape() as t2: - with tf.GradientTape(persistent=True) as t1: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], dev, gradient_fn=param_shift, interface=interface, max_diff=2 - )[0] - res = tf.stack(res) - - g = t1.jacobian(res, params, experimental_use_pfor=False) - - hess = t2.jacobian(g, params) - - a, b = params * 1.0 + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost_fn(params) + grad = grad_tape.gradient(res, params) + hess = jac_tape.jacobian(grad, params) - expected_res = [ - 0.5 + 0.5 * tf.cos(a) * tf.cos(b), - 0.5 - 0.5 * tf.cos(a) * tf.cos(b), - ] - assert np.allclose(res, expected_res, atol=tol, rtol=0) + x, y = params + expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - expected_g = [ - [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)], - [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)], - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) + expected = tf.convert_to_tensor( + [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) - expected_hess = [ - [ - [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)], - [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)], - ], + expected = tf.convert_to_tensor( [ - [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)], - [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)], - ], - ] - - np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - - def test_adjoint_hessian(self, interface): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface=interface, - )[0] - - grad = t1.gradient(res, params) - assert grad is not None - assert grad.dtype == tf.float64 - assert grad.shape == params.shape - - hess = t2.jacobian(grad, params) - assert hess is None + [-tf.cos(2 * x) * tf.cos(2 * y), tf.sin(2 * x) * tf.sin(2 * y)], + [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], + ] + ) + assert np.allclose(hess, expected, atol=tol, rtol=0) - def test_max_diff(self, tol, interface): + def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.tf", wires=2) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - x, y = params * 1.0 - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(params[0], wires=0) - qml.RY(params[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface=interface - ) - res = result[0] + result[1][0] + dev = DefaultQubit() + params = tf.Variable([0.543, -0.654]) - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost_fn(x): + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - grad = t1.gradient(res, params) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] - hess = t2.jacobian(grad, params) - assert hess is None + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost_fn(params) + grad = grad_tape.gradient(res, params) + hess = jac_tape.gradient(grad, params) + x, y = params + expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift, "interface": "tensorflow"}, - {"gradient_fn": finite_diff, "interface": "tensorflow"}, - {"gradient_fn": param_shift, "interface": "auto"}, - {"gradient_fn": finite_diff, "interface": "auto"}, -] + expected = tf.convert_to_tensor( + [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + assert hess is None -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -956,73 +734,111 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - return tf.stack(execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") return _cost_fn @staticmethod def cost_fn_expected(weights, coeffs1, coeffs2): """Analytic value of cost_fn above""" - a, b, c = coeffs1.numpy() - d = coeffs2.numpy()[0] - x, y = weights.numpy() - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * tf.sin(x) * tf.sin(y) + tf.cos(x) * (a + b * tf.sin(y)), d * tf.cos(x)] @staticmethod - def cost_fn_jacobian_expected(weights, coeffs1, coeffs2): + def cost_fn_jacobian(weights, coeffs1, coeffs2): """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.numpy() - d = coeffs2.numpy()[0] - x, y = weights.numpy() - return np.array( + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return tf.convert_to_tensor( [ [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), + -c * tf.cos(x) * tf.sin(y) - tf.sin(x) * (a + b * tf.sin(y)), + b * tf.cos(x) * tf.cos(y) - c * tf.cos(y) * tf.sin(x), + tf.cos(x), + tf.cos(x) * tf.sin(y), + -(tf.sin(x) * tf.sin(y)), 0, ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + [-d * tf.sin(x), 0, 0, 0, 0, tf.cos(x)], ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + device_vjp = execute_kwargs.get("device_vjp", False) + coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) coeffs2 = tf.constant([0.7], dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = tape.jacobian(res, [weights, coeffs1, coeffs2]) - expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) - assert res[1] is None - assert res[2] is None + jac = tape.jacobian(res, [weights], experimental_use_pfor=not device_vjp) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) coeffs2 = tf.Variable([0.7], dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac = qml.math.hstack(tape.jacobian(res, [weights, coeffs1, coeffs2]), like="tensorflow") + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) + + +@pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) +def test_device_returns_float32(diff_method): + """Test that if the device returns float32, the derivative succeeds.""" + + def _to_float32(results): + if isinstance(results, (list, tuple)): + return tuple(_to_float32(r) for r in results) + return np.array(results, dtype=np.float32) + + class Float32Dev(qml.devices.DefaultQubit): + def execute(self, circuits, execution_config=qml.devices.DefaultExecutionConfig): + results = super().execute(circuits, execution_config) + return _to_float32(results) + + dev = Float32Dev() + + @qml.qnode(dev, diff_method=diff_method) + def circuit(x): + qml.RX(tf.cos(x), wires=0) + return qml.expval(qml.Z(0)) + + x = tf.Variable(0.1, dtype=tf.float64) + + with tf.GradientTape() as tape: + y = circuit(x) + + assert qml.math.allclose(y, np.cos(np.cos(0.1))) - res = tape.jacobian(res, [weights, coeffs1, coeffs2]) - expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) - assert np.allclose(res[1], expected[:, 2:5], atol=tol, rtol=0) - assert np.allclose(res[2], expected[:, 5:], atol=tol, rtol=0) + g = tape.gradient(y, x) + expected_g = np.sin(np.cos(0.1)) * np.sin(0.1) + assert qml.math.allclose(g, expected_g) diff --git a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py index ab3147dfc95..695438310c1 100644 --- a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py +++ b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods +# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg,redefined-outer-name import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -26,22 +27,35 @@ shots_and_num_copies = [((1, (5, 2), 10), 4)] shots_and_num_copies_hess = [((10, (5, 1)), 2)] -spsa_kwargs = {"h": 10e-2, "num_directions": 30, "sampler_rng": np.random.default_rng(42)} +kwargs = { + "finite-diff": {"h": 10e-2}, + "parameter-shift": {}, + "spsa": {"h": 10e-2, "num_directions": 20}, +} + qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", spsa_kwargs], + [DefaultQubit(), "finite-diff"], + [DefaultQubit(), "parameter-shift"], + [DefaultQubit(), "spsa"], ] TOLS = { "finite-diff": 0.3, "parameter-shift": 1e-2, - "spsa": 0.3, + "spsa": 0.5, } +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -50,14 +64,13 @@ class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and one param, the gradient is a float.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)) @@ -65,7 +78,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,14 +87,13 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)) @@ -90,7 +102,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -102,14 +114,13 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) @@ -117,7 +128,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -126,15 +137,14 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.probs(wires=[0, 1]) @@ -142,7 +152,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -151,15 +161,14 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) @@ -168,7 +177,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -180,15 +189,14 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) @@ -196,7 +204,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -205,24 +213,22 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=1, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -234,14 +240,13 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) qml.RY(a[2], wires=0) @@ -250,7 +255,7 @@ def circuit(a): a = tf.Variable([0.7, 0.9, 1.1], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -259,14 +264,13 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a single params return an array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -274,7 +278,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -283,14 +287,13 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -299,7 +302,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -311,14 +314,13 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -326,7 +328,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -337,7 +339,7 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -346,7 +348,7 @@ class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" @@ -354,14 +356,12 @@ def test_hessian_expval_multiple_params( # TODO: Find out why. pytest.skip("SPSA gradient does not support this particular test case") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) @@ -369,7 +369,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -384,11 +384,11 @@ def circuit(x, y): assert h.shape == (2, num_copies) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -397,24 +397,23 @@ class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -431,24 +430,23 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/test_tensorflow_qnode.py b/tests/interfaces/test_tensorflow_qnode.py index 5b03f9da10f..e3c4597ae06 100644 --- a/tests/interfaces/test_tensorflow_qnode.py +++ b/tests/interfaces/test_tensorflow_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2020 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,26 +13,31 @@ # limitations under the License. """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, use-implicit-booleaness-not-comparison import pytest import pennylane as qml from pennylane import qnode - -# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison - +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf tf = pytest.importorskip("tensorflow") - +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=4), "adjoint", False, True], + [qml.device("lightning.qubit", wires=4), "adjoint", False, False], + [qml.device("lightning.qubit", wires=4), "adjoint", True, True], + [qml.device("lightning.qubit", wires=4), "adjoint", True, False], ] TOL_FOR_SPSA = 1.0 @@ -45,25 +50,25 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with TensorFlow integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -84,7 +89,7 @@ def circuit(a): # with the interface, the tape returns tensorflow tensors assert isinstance(res, tf.Tensor) - assert res.shape == tuple() + assert res.shape == () # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -92,23 +97,20 @@ def circuit(a): # gradients should work grad = tape.gradient(res, a) assert isinstance(grad, tf.Tensor) - assert grad.shape == tuple() + assert grad.shape == () - def test_interface_swap(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_interface_swap(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test that the TF interface can be applied to a QNode with a pre-existing interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -135,51 +137,44 @@ def circuit(a): assert np.allclose(res1, res2, atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, dev_name, diff_method, grad_on_execution, interface): + def test_drawing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test circuit drawing when using the TF interface""" x = tf.Variable(0.1, dtype=tf.float64) y = tf.Variable([0.2, 0.3], dtype=tf.float64) z = tf.Variable(0.4, dtype=tf.float64) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) qml.RY(p2[0] * p2[1], wires=1) qml.RX(kwargs["p3"], wires=0) qml.CNOT(wires=[0, 1]) - return qml.state() + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) result = qml.draw(circuit)(p1=x, p3=z) - expected = "0: ──RX(0.10)──RX(0.40)─╭●── State\n1: ──RY(0.06)───────────╰X── State" + expected = "0: ──RX(0.10)──RX(0.40)─╭●── \n1: ──RY(0.06)───────────╰X── " assert result == expected - def test_jacobian(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_jacobian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test jacobian calculation""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) @@ -190,7 +185,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -202,61 +197,17 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [a, b]) + res = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) expected = [[-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_dtype(self, dev_name, diff_method, grad_on_execution, interface): - """Test calculating the jacobian with a different datatype""" - if diff_method == "backprop": - pytest.skip("Test does not support backprop") - - a = tf.Variable(0.1, dtype=tf.float32) - b = tf.Variable(0.2, dtype=tf.float32) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, r_dtype=np.float32) - - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - - with tf.GradientTape() as tape: - res = circuit(a, b) - res = tf.stack(res) - - assert circuit.qtape.trainable_params == [0, 1] - - assert isinstance(res, tf.Tensor) - assert res.shape == (2,) - assert res.dtype is tf.float32 - - res = tape.jacobian(res, [a, b]) - assert [r.dtype is tf.float32 for r in res] - - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test setting finite-difference jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite diff and spsa.") a = tf.Variable([0.1, 0.2]) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, interface=interface, @@ -264,18 +215,21 @@ def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interf approx_order=2, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - tape.jacobian(res, a) + tape.jacobian(res, a, experimental_use_pfor=not device_vjp) - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method in ["backprop", "adjoint", "spsa"]: @@ -284,21 +238,16 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, t a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) - num_wires = 2 - diff_kwargs = {} - if diff_method == "hadamard": - num_wires = 3 - elif diff_method == "finite-diff": + if diff_method == "finite-diff": diff_kwargs = {"approx_order": 2, "strategy": "center"} - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **diff_kwargs, ) def circuit(a, b): @@ -307,7 +256,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -317,7 +266,7 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, [a, b]) + jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) expected = [ [-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)], @@ -328,7 +277,7 @@ def circuit(a, b): a = tf.Variable(0.54, dtype=tf.float64) b = tf.constant(0.8, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -338,25 +287,22 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) expected = [-tf.sin(a), tf.sin(a) * tf.sin(b)] assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(x, y, z): qml.RY(x * z, wires=0) @@ -364,30 +310,30 @@ def circuit(x, y, z): qml.RX(z + z**2 + tf.sin(a), wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b, c) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [0, 2] assert circuit.qtape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] - res = tape.jacobian(res, [a, b, c]) + res = tape.jacobian(res, [a, b, c], experimental_use_pfor=not device_vjp) assert isinstance(res[0], tf.Tensor) assert res[1] is None assert isinstance(res[2], tf.Tensor) - def test_no_trainable_parameters(self, dev_name, diff_method, grad_on_execution, interface): + def test_no_trainable_parameters( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """Test evaluation if there are no trainable parameters""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -398,7 +344,7 @@ def circuit(a, b): a = 0.1 b = tf.constant(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -409,31 +355,30 @@ def circuit(a, b): assert isinstance(res, tf.Tensor) # can't take the gradient with respect to "a" since it's a Python scalar - grad = tape.jacobian(res, b) + grad = tape.jacobian(res, b, experimental_use_pfor=not device_vjp) assert grad is None @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, U, tol, interface): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, device_vjp, U, tol, interface + ): """Test that the TF interface works correctly with a matrix parameter""" a = tf.Variable(0.1, dtype=tf.float64) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) qml.RY(a, wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(U, a) if diff_method == "finite-diff": @@ -441,18 +386,23 @@ def circuit(U, a): assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) - res = tape.jacobian(res, a) + res = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test that operation and nested tapes expansion is differentiable""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA class U3(qml.U3): @@ -464,13 +414,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - a = np.array(0.1) p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) @@ -480,7 +423,7 @@ def circuit(a, p): U3(p[0], p[1], p[2], wires=0) return qml.expval(qml.PauliX(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( @@ -488,7 +431,7 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, p) + res = tape.jacobian(res, p, experimental_use_pfor=not device_vjp) expected = np.array( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), @@ -507,9 +450,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, mocker, tol, interface): + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -518,31 +461,21 @@ def circuit(weights): qml.RY(weights[0], wires=0) qml.RX(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(weights) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(weights) # execute with shots=100 - circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(weights) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() + res = circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg + assert res.shape == (100, 2) def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" # pylint: disable=unexpected-keyword-arg - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -557,7 +490,6 @@ def circuit(weights): res = circuit(weights, shots=[10000, 10000, 10000]) res = tf.transpose(tf.stack(res)) - assert dev.shots is None assert len(res) == 3 jacobian = tape.jacobian(res, weights) @@ -568,7 +500,7 @@ def test_multiple_gradient_integration(self, tol, interface): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -582,7 +514,7 @@ def circuit(weights): with tf.GradientTape() as tape: res1 = circuit(weights) - assert qml.math.shape(res1) == tuple() + assert qml.math.shape(res1) == () res2 = circuit(weights, shots=[(1, 1000)]) # pylint: disable=unexpected-keyword-arg assert qml.math.shape(res2) == (1000,) @@ -593,7 +525,7 @@ def circuit(weights): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) + dev = DefaultQubit() weights = tf.Variable([0.543, -0.654], dtype=tf.float64) spy = mocker.spy(qml, "execute") @@ -605,115 +537,44 @@ def circuit(weights): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + circuit(weights, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - circuit(weights) - assert circuit.gradient_fn is qml.gradients.param_shift - - circuit(weights) + assert circuit.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - assert circuit.gradient_fn is qml.gradients.param_shift - - # if we set the shots to None, backprop can now be used - circuit(weights, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert circuit.gradient_fn == "backprop" + # if we use the default shots value of None, backprop can now be used circuit(weights) - assert circuit.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - - -@pytest.mark.parametrize("interface", ["auto", "tf"]) -class TestAdjoint: - """Specific integration tests for the adjoint method""" - - def test_reuse_state(self, mocker, interface): - """Tests that the TF interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, diff_method="adjoint", interface=interface) - def circ(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=1) - qml.CNOT(wires=(0, 1)) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - weights = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = 1.0 * weights - - with tf.GradientTape() as tape: - res = tf.reduce_sum(circ(weights)) - - grad = tape.gradient(res, weights) - expected_grad = [-tf.sin(x), tf.cos(y)] - - assert np.allclose(grad, expected_grad) - assert circ.device.num_executions == 1 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) - - def test_reuse_state_multiple_evals(self, mocker, tol, interface): - """Tests that the TF interface reuses the device state for adjoint differentiation, - even where there are intermediate evaluations.""" - dev = qml.device("default.qubit.legacy", wires=2) - - x_val = 0.543 - y_val = -0.654 - x = tf.Variable(x_val, dtype=tf.float64) - y = tf.Variable(y_val, dtype=tf.float64) - - @qnode(dev, diff_method="adjoint", interface=interface) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - with tf.GradientTape() as tape: - res1 = circuit(x, y) - - assert np.allclose(res1, np.cos(x_val), atol=tol, rtol=0) - - # intermediate evaluation with different values - circuit(tf.math.tan(x), tf.math.cosh(y)) - - # the adjoint method will continue to compute the correct derivative - grad = tape.gradient(res1, x) - assert np.allclose(grad, -np.sin(x_val), atol=tol, rtol=0) - assert dev.num_executions == 2 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + assert circuit.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, dev_name, diff_method, grad_on_execution, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface ): """Tests correct output shape and evaluation for a tape with multiple probs outputs""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - num_wires = 2 + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 + tol = TOL_FOR_SPSA x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -725,7 +586,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.probs(wires=[0]), qml.probs(wires=[1]) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(x, y) res = tf.stack(res) @@ -737,7 +598,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y]) + res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) expected = np.array( [ [ @@ -752,24 +613,25 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_ragged_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - num_wires = 2 + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 + tol = TOL_FOR_SPSA x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -781,7 +643,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(x, y) res = tf.experimental.numpy.hstack(res) @@ -794,7 +656,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y]) + res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) expected = np.array( [ [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], @@ -803,24 +665,20 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -850,24 +708,18 @@ def circuit(x): ] assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -896,24 +748,20 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -956,25 +804,19 @@ def circuit(x): np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, grad_on_execution, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=0) @@ -1013,24 +855,18 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian_ragged(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -1081,23 +917,21 @@ def circuit(x): ] np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - def test_state(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1107,7 +941,6 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is tf.complex128 probs = tf.math.abs(res) ** 2 return probs[0] + probs[2] @@ -1122,21 +955,27 @@ 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 - def test_projector(self, state, dev_name, diff_method, grad_on_execution, tol, interface): + @pytest.mark.parametrize("dtype", ("int32", "int64")) + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, tol, interface, dtype + ): """Test that the variance of a projector is correctly returned""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all expvals or all diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Variance not implemented yet.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) - P = tf.constant(state) + P = tf.constant(state, dtype=dtype) x, y = 0.765, -0.654 weights = tf.Variable([x, y], dtype=tf.float64) @@ -1161,95 +1000,69 @@ def circuit(weights): ] assert np.allclose(grad, expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = tf.Variable(0.543, dtype=tf.float64) - phi = tf.Variable(-0.654, dtype=tf.float64) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - with tf.GradientTape() as tape: - res = circuit(r, phi) - - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - grad = tape.gradient(res, [r, phi]) - expected = [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - assert np.allclose(grad, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - n = tf.Variable(0.12, dtype=tf.float64) - a = tf.Variable(0.765, dtype=tf.float64) + phi = tf.Variable(1.23) + theta = tf.Variable(4.56) - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - with tf.GradientTape() as tape: - res = circuit(n, a) + with tf.GradientTape() as res_tape: + res = circuit(phi, theta) + gradient = res_tape.gradient(res, [phi, theta]) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) + with tf.GradientTape() as expected_tape: + expected = expected_circuit(theta) + exp_theta_grad = expected_tape.gradient(expected, theta) - # circuit jacobians - grad = tape.gradient(res, [n, a]) - expected = [2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + assert np.allclose(gradient, [0.0, exp_theta_grad]) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the TF interface""" - def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution, interface): + def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) class PhaseShift(qml.PhaseShift): grad_method = None @@ -1263,6 +1076,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.Hadamard(wires=0) @@ -1274,7 +1088,6 @@ def circuit(x): with tf.GradientTape() as t2: with tf.GradientTape() as t1: loss = circuit(x) - res = t1.gradient(loss, x) assert np.allclose(res, -3 * np.sin(3 * x)) @@ -1286,20 +1099,13 @@ def circuit(x): @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - class PhaseShift(qml.PhaseShift): grad_method = None @@ -1312,6 +1118,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1329,24 +1136,24 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, tol, interface ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" - kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - interface=interface, - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "max_diff": max_diff, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint/hadamard method does not yet support Hamiltonians") elif diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1393,14 +1200,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker, interface + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1413,8 +1220,6 @@ def test_hamiltonian_expansion_finite_shots( elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1423,6 +1228,7 @@ def test_hamiltonian_expansion_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1438,17 +1244,19 @@ def circuit(data, weights, coeffs): c = tf.Variable([-0.6543, 0.24, 0.54], dtype=tf.float64) # test output - with tf.GradientTape(persistent=True) as t2: + with tf.GradientTape(persistent=True) as _t2: with tf.GradientTape() as t1: - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients grad = t1.gradient(res, [d, w, c]) + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1458,33 +1266,36 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[2], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - grad2_c = t2.jacobian(grad[2], c) - assert grad2_c is None + # TODO: figure out why grad2_c is np.zeros((3,3)) instead of None + # if diff_method == "parameter-shift" and max_diff == 2: + # grad2_c = _t2.jacobian(grad[2], c) + # print(grad2_c, grad[2], c) + # assert grad2_c is None - grad2_w_c = t2.jacobian(grad[1], c) - expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ - 0, - -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), - -np.sin(d[1] + w[1]), - ] - assert np.allclose(grad2_w_c, expected, atol=tol) + # grad2_w_c = _t2.jacobian(grad[1], c) + # expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ + # 0, + # -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), + # -np.sin(d[1] + w[1]), + # ] + # assert np.allclose(grad2_w_c, expected, atol=tol) class TestSample: """Tests for the sample integration""" + # pylint:disable=unexpected-keyword-arg + def test_sample_dimension(self): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert len(res) == 2 @@ -1496,15 +1307,14 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) @@ -1515,76 +1325,67 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" - n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=10) assert len(result) == 3 - assert result[0].shape == (n_sample,) + assert result[0].shape == (10,) assert result[1].shape == () assert result[2].shape == () assert isinstance(result[0], tf.Tensor) assert isinstance(result[1], tf.Tensor) assert isinstance(result[2], tf.Tensor) - assert result[0].dtype is tf.float64 - assert result[1].dtype is tf.float64 - assert result[2].dtype is tf.float64 + assert result[0].dtype is tf.float64 # pylint:disable=no-member + assert result[1].dtype is tf.float64 # pylint:disable=no-member + assert result[2].dtype is tf.float64 # pylint:disable=no-member def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (n_sample,)) + assert np.array_equal(result.shape, (10,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=10) result = tf.stack(result) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (3, n_sample)) + assert np.array_equal(result.shape, (3, 10)) assert result.dtype == tf.float64 def test_counts(self): """Test counts works as expected for TF""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) - @qnode(dev, interface="tf") + # pylint:disable=unsubscriptable-object,no-member + @qnode(DefaultQubit(), interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)) - res = circuit() + res = circuit(shots=100) assert isinstance(res, dict) assert list(res.keys()) == [-1, 1] @@ -1614,12 +1415,11 @@ class TestAutograph: def test_autograph_gradients(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1628,7 +1428,7 @@ def circuit(x, y): with tf.GradientTape() as tape: p0, p1 = circuit(x, y) - loss = p0[0] + p1[1] + loss = p0[0] + p1[1] # pylint:disable=unsubscriptable-object expected = tf.cos(x / 2) ** 2 + (1 - tf.cos(x) * tf.cos(y)) / 2 assert np.allclose(loss, expected, atol=tol, rtol=0) @@ -1640,12 +1440,11 @@ def circuit(x, y): def test_autograph_jacobian(self, decorator, interface, tol): """Test that a parameter-shift vector-valued QNode can be compiled using @tf.function, and differentiated""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1683,10 +1482,14 @@ def circuit(x, y): def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) def circuit(x): qml.RY(x, wires=0) return qml.expval(qml.PauliZ(0)) @@ -1707,10 +1510,14 @@ def circuit(x): def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -1729,16 +1536,44 @@ def circuit(x): expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)] assert np.allclose(g, expected_g, atol=tol, rtol=0) + @pytest.mark.xfail + @pytest.mark.parametrize("grad_on_execution", [True, False]) + def test_autograph_adjoint_multi_out(self, grad_on_execution, decorator, interface, tol): + """Test that a parameter-shift QNode can be compiled + using @tf.function, and differentiated to second order""" + + @decorator + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) + + x = tf.Variable(0.5, dtype=tf.float64) + + with tf.GradientTape() as tape: + res = qml.math.hstack(circuit(x)) + g = tape.jacobian(res, x) + a = x * 1.0 + + expected_res = [tf.cos(a), tf.sin(a)] + assert np.allclose(res, expected_res, atol=tol, rtol=0) + + expected_g = [-tf.sin(a), tf.cos(a)] + assert np.allclose(g, expected_g, atol=tol, rtol=0) + def test_autograph_ragged_differentiation(self, decorator, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device("default.qubit.legacy", wires=2) - x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1770,12 +1605,11 @@ def circuit(x, y): def test_autograph_hessian(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) a = tf.Variable(0.543, dtype=tf.float64) b = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=2, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=2, interface=interface) def circuit(x, y): qml.RY(x, wires=0) qml.RX(y, wires=0) @@ -1801,35 +1635,15 @@ def circuit(x, y): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_autograph_sample(self, decorator, interface): - """Test that a QNode returning raw samples can be compiled - using @tf.function""" - dev = qml.device("default.qubit", wires=2, shots=100) - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.sample() - - result = circuit(x, y) - result = np.array(result).flatten() - assert np.all([r in [1, 0] for r in result]) - def test_autograph_state(self, decorator, interface, tol): """Test that a parameter-shift QNode returning a state can be compiled using @tf.function""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) # TODO: fix this for diff_method=None @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1846,16 +1660,15 @@ def circuit(x, y): def test_autograph_dimension(self, decorator, interface): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) - def circuit(): + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + def circuit(**_): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) # pylint:disable=unexpected-keyword-arg res = tf.stack(res) assert res.shape == (2, 10) @@ -1863,24 +1676,23 @@ def circuit(): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and one param, the gradient is a float.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1898,19 +1710,16 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1931,19 +1740,16 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1962,68 +1768,64 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) + a = tf.Variable(0.1, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) - b = tf.Variable(0.2) + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) - jac = tape.jacobian(res, (a, b)) + jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) assert isinstance(jac, tuple) @@ -2034,101 +1836,94 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single array is returned.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2]) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a single params return an array.""" - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) + a = tf.Variable(0.1, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) - b = tf.Variable(0.2) + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, (a, b)) + jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2140,50 +1935,40 @@ def circuit(a, b): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2]) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (5, 2) def test_hessian_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2196,6 +1981,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2222,20 +2008,13 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2244,6 +2023,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2262,10 +2042,10 @@ def circuit(x): assert isinstance(hess, tf.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, interface): + def test_hessian_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2281,6 +2061,7 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2307,7 +2088,7 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2316,8 +2097,6 @@ def test_hessian_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") - dev = qml.device(dev_name, wires=2) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2326,6 +2105,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2345,16 +2125,9 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - pytest.skip("Test does not support hadamard because multiple measurements.") - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2367,6 +2140,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2394,7 +2168,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" @@ -2404,10 +2178,6 @@ def test_hessian_probs_expval_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because multiple measurements.") - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2416,6 +2186,7 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2436,11 +2207,9 @@ def circuit(x): assert hess.shape == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") if diff_method == "hadamard": @@ -2455,6 +2224,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2482,7 +2252,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2490,8 +2260,6 @@ def test_hessian_probs_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") - dev = qml.device(dev_name, wires=2) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2500,6 +2268,7 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2520,17 +2289,49 @@ def circuit(x): assert hess.shape == (3, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="tf") + @qml.qnode(DefaultQubit(), interface="tf") def circuit(): qml.Hadamard(wires=0) return qml.state() res = circuit() assert isinstance(res, tf.Tensor) + + +def test_error_device_vjp_jacobian(): + """Test a ValueError is raised if a jacobian is attempted to be computed with device_vjp=True.""" + + dev = qml.device("default.qubit") + + @qml.qnode(dev, diff_method="adjoint", device_vjp=True) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) + + x = tf.Variable(0.1) + + with tf.GradientTape() as tape: + y = qml.math.hstack(circuit(x)) + + with pytest.raises(ValueError): + tape.jacobian(y, x) + + +def test_error_device_vjp_state_float32(): + """Test a ValueError is raised is state differentiation is attemped with float32 parameters.""" + + dev = qml.device("default.qubit") + + @qml.qnode(dev, diff_method="adjoint", device_vjp=True) + def circuit(x): + qml.RX(x, wires=0) + return qml.probs(wires=0) + + x = tf.Variable(0.1, dtype=tf.float32) + with pytest.raises(ValueError, match="tensorflow with adjoint differentiation of the state"): + with tf.GradientTape(): + circuit(x) diff --git a/tests/interfaces/test_tensorflow_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_qnode_shot_vector.py index 614c71ee7d4..c0340aaab9a 100644 --- a/tests/interfaces/test_tensorflow_qnode_shot_vector.py +++ b/tests/interfaces/test_tensorflow_qnode_shot_vector.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments +# pylint: disable=too-many-arguments,unexpected-keyword-arg import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -27,9 +28,9 @@ shots_and_num_copies_hess = [((10, (5, 1)), 2)] qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], + [DefaultQubit(), "finite-diff", {"h": 10e-2}], + [DefaultQubit(), "parameter-shift", {}], + [DefaultQubit(), "spsa", {"h": 10e-2, "num_directions": 20}], ] TOLS = { @@ -45,16 +46,15 @@ @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and one param, the gradient is a float.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -65,7 +65,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,10 +74,9 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -89,7 +88,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -101,10 +100,9 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -115,7 +113,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -124,11 +122,10 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -139,7 +136,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -148,11 +145,10 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -164,7 +160,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -176,11 +172,10 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -191,7 +186,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -200,10 +195,9 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1) par_1 = tf.Variable(0.2) @@ -216,7 +210,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -228,10 +222,9 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -243,7 +236,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2, 0.3]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -252,10 +245,9 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a single params return an array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -266,7 +258,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -275,10 +267,9 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -290,7 +281,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -302,10 +293,9 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -316,7 +306,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -328,16 +318,15 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -351,7 +340,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -366,10 +355,9 @@ def circuit(x, y): assert h.shape == (2, num_copies) def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) params = tf.Variable([0.1, 0.2], dtype=tf.float64) @@ -382,7 +370,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params) + res = circuit(params, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -393,10 +381,9 @@ def circuit(x): assert hess.shape == (num_copies, 2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -410,7 +397,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -425,12 +412,10 @@ def circuit(x, y): assert h.shape == (2, num_copies, 3) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) @@ -442,7 +427,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params) + res = circuit(params, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -453,22 +438,22 @@ def circuit(x): assert hess.shape == (num_copies, 3, 2, 2) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - dev = qml.device(dev_name, wires=2, shots=shots) + x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -480,7 +465,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -497,11 +482,11 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device(dev_name, wires=2, shots=shots) + x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -513,7 +498,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py index 5aac97fb0e3..3cdcf5eae30 100644 --- a/tests/interfaces/test_torch.py +++ b/tests/interfaces/test_torch.py @@ -1,4 +1,4 @@ -# Copyright 2018-2022 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,254 +11,45 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the Torch interface""" -# pylint: disable=protected-access,too-few-public-methods +"""Torch specific tests for execute and default qubit 2.""" import numpy as np import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import execute -from pennylane.gradients import finite_diff, param_shift -from pennylane.typing import TensorLike +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift +from pennylane.measurements import Shots -pytestmark = pytest.mark.torch torch = pytest.importorskip("torch") -torch_functional = pytest.importorskip("torch.autograd.functional") -torch_cuda = pytest.importorskip("torch.cuda") - - -@pytest.mark.parametrize("interface", ["torch", "auto"]) -class TestTorchExecuteUnitTests: - """Unit tests for torch execution""" - - def test_jacobian_options(self, interface, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute( - [tape], - dev, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - interface=interface, - )[0] - - res.backward() - - for args in spy.call_args_list: - assert args[1]["shift"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self, interface): - """Test that an error is raised if a gradient transform - is used with grad_on_execution=True""" - a = torch.tensor([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - execute( - [tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface=interface - ) - - def test_grad_on_execution_reuse_state(self, interface, mocker): - """Test that grad_on_execution uses the `device.execute_and_gradients` pathway - while reusing the quantum state.""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface=interface, - ) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_grad_on_execution(self, interface, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian"}, - interface=interface, - ) - - # two device executions; one for the value, one for the Jacobian - assert dev.num_executions == 2 - spy.assert_called() - - def test_no_grad_on_execution(self, interface, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface=interface, - )[0] +pytestmark = pytest.mark.torch - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - res.backward() - spy_gradients.assert_called() +@pytest.fixture(autouse=True) +def run_before_and_after_tests(): + torch.set_default_dtype(torch.float64) + yield + torch.set_default_dtype(torch.float32) +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute( - [tape], dev, gradient_fn=param_shift, cachesize=cachesize, interface="torch" - )[0][0] - - params = torch.tensor([0.1, 0.2], requires_grad=True) - res = cost(params, cachesize=2) - res.backward() - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ - 0 - ] - - custom_cache = {} - params = torch.tensor([0.1, 0.2], requires_grad=True) - res = cost(params, cache=custom_cache) - res.backward() - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_caching_param_shift(self): - """Test that, with the parameter-shift transform, - Torch always uses the optimum number of evals when computing the Jacobian.""" - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ - 0 - ] - - # Without caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - params = torch.tensor([0.1, 0.2], requires_grad=True) - torch_functional.jacobian(lambda p: cost(p, cache=None), params) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - torch_functional.jacobian(lambda p: cost(p, cache=True), params) - assert dev.num_executions == 5 + """Tests for caching behaviour""" + @pytest.mark.skip("caching is not implemented for torch") @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): - """Test that, with the parameter-shift transform, + def test_caching_param_shift_hessian(self, num_params): + """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor(np.arange(1, num_params + 1) / 10, requires_grad=True) + dev = DefaultQubit() + params = torch.arange(1, num_params + 1, requires_grad=True, dtype=torch.float64) N = len(params) - def cost(x, cache): + def get_cost_tape(x): with qml.queuing.AnnotatedQueue() as q: qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -267,447 +58,447 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - tape = qml.tape.QuantumScript.from_queue(q) + return qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], dev, gradient_fn=param_shift, cache=cache, interface="torch", max_diff=2 + def cost_no_cache(x): + return qml.execute( + [get_cost_tape(x)], + dev, + gradient_fn=qml.gradients.param_shift, + cache=False, + max_diff=2, + )[0] + + def cost_cache(x): + return qml.execute( + [get_cost_tape(x)], + dev, + gradient_fn=qml.gradients.param_shift, + cache=True, + max_diff=2, )[0] # No caching: number of executions is not ideal - hess1 = torch.autograd.functional.hessian(lambda x: cost(x, cache=None), params) + with qml.Tracker(dev) as tracker: + hess1 = torch.autograd.functional.hessian(cost_no_cache, params) if num_params == 2: # compare to theoretical result - x, y, *_ = params.detach() + x, y, *_ = params.clone().detach() expected = torch.tensor( [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [2 * torch.cos(2 * x) * torch.sin(y) ** 2, torch.sin(2 * x) * torch.sin(2 * y)], + [ + torch.sin(2 * x) * torch.sin(2 * y), + -2 * torch.cos(x) ** 2 * torch.cos(2 * y), + ], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert torch.allclose(expected, hess1) expected_runs = 1 # forward pass - expected_runs += 2 * N # Jacobian - expected_runs += 4 * N + 1 # Hessian diagonal - expected_runs += 4 * N**2 # Hessian off-diagonal - assert dev.num_executions == expected_runs + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - hess2 = torch.autograd.functional.hessian(lambda x: cost(x, cache=True), params) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + with qml.Tracker(dev) as tracker2: + hess2 = torch.autograd.functional.hessian(cost_cache, params) + assert torch.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal + assert tracker2.totals["executions"] == expected_runs_ideal assert expected_runs_ideal < expected_runs - def test_caching_adjoint_no_grad_on_execution(self): - """Test that caching reduces the number of adjoint evaluations - when grad_on_execution=False""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.1, 0.2, 0.3]) - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), + ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, + Shots(None), + DefaultQubit(), + ), + ( + { + "gradient_fn": "adjoint", + "grad_on_execution": False, + "device_vjp": False, + }, + Shots(None), + DefaultQubit(), + ), + ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "device", "device_vjp": False}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(), + ), + ( + {"gradient_fn": "device", "device_vjp": True}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(), + ), +] - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface="torch", - )[0] - - # Without caching, 2 evaluations are required. - # 1 for the forward pass, and one for the backward pass - torch_functional.jacobian(lambda x: cost(x, cache=None), params) - assert dev.num_executions == 2 - - # With caching, only 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - torch_functional.jacobian(lambda x: cost(x, cache=True), params) - assert dev.num_executions == 2 - - -torch_devices = [None] - -if torch_cuda.is_available(): - torch_devices.append(torch.device("cuda")) - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift, "interface": "torch"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, - "interface": "torch", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "torch", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "torch", - }, - {"gradient_fn": param_shift, "interface": "auto"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "auto", - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 -@pytest.mark.gpu -@pytest.mark.parametrize("torch_device", torch_devices) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestTorchExecuteIntegration: """Test the torch interface execute function integrates well for both forward and backward execution""" - def test_execution(self, torch_device, execute_kwargs): - """Test that the execute function produces results with the expected shapes""" - dev = qml.device("default.qubit.legacy", wires=1) - a = torch.tensor(0.1, requires_grad=True, device=torch_device) - b = torch.tensor(0.2, requires_grad=False, device=torch_device) - - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) + def test_execution(self, execute_kwargs, shots, device): + """Test execution""" - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b): + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - res = execute([tape1, tape2], dev, **execute_kwargs) + return execute([tape1, tape2], device, **execute_kwargs) - assert isinstance(res, TensorLike) - assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=False) - def test_scalar_jacobian(self, torch_device, execute_kwargs, tol): - """Test scalar jacobian calculation by comparing two types of pipelines""" - a = torch.tensor(0.1, requires_grad=True, dtype=torch.float64, device=torch_device) - dev = qml.device("default.qubit.legacy", wires=2) + with device.tracker: + res = cost(a, b) - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + if execute_kwargs.get("grad_on_execution", False): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes - tape = qml.tape.QuantumScript.from_queue(q) + assert len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () + exp = torch.cos(a) * torch.cos(b) + if shots.has_partitioned_shots: + for shot in range(2): + for wire in range(2): + assert qml.math.allclose(res[shot][wire], exp, atol=atol_for_shots(shots)) + else: + for wire in range(2): + assert qml.math.allclose(res[wire], exp, atol=atol_for_shots(shots)) - res = execute([tape], dev, **execute_kwargs)[0] - res.backward() + def test_scalar_jacobian(self, execute_kwargs, shots, device): + """Test scalar jacobian calculation""" + a = torch.tensor(0.1, requires_grad=True) - # compare to backprop gradient def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - dev = qml.device("default.qubit.autograd", wires=2) - return dev.batch_execute([tape])[0] - - expected = qml.grad(cost, argnum=0)(0.1) - assert torch.allclose(a.grad, torch.tensor(expected, device=torch_device), atol=tol, rtol=0) - - def test_jacobian(self, torch_device, execute_kwargs, tol): - """Test jacobian calculation by checking against analytic values""" - a_val = 0.1 - b_val = 0.2 - - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - b = torch.tensor(b_val, requires_grad=True, device=torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RZ(torch.tensor(0.543, device=torch_device), wires=0) - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1, 2] - - assert isinstance(res, tuple) - - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == () + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] + + res = torch.autograd.functional.jacobian(cost, a) + if not shots.has_partitioned_shots: + assert res.shape == () # pylint: disable=no-member + + # compare to standard tape jacobian + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(device.execute(tapes)) + + assert expected.shape == () + if shots.has_partitioned_shots: + for i in range(shots.num_copies): + assert torch.allclose(res[i], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[i], -torch.sin(a), atol=atol_for_shots(shots)) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res, -torch.sin(a), atol=atol_for_shots(shots)) + + def test_jacobian(self, execute_kwargs, shots, device): + """Test jacobian calculation""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=True) + + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + [res] = execute([tape], device, **execute_kwargs) + if shots.has_partitioned_shots: + return torch.hstack(res[0] + res[1]) + return torch.hstack(res) + + res = cost(a, b) + expected = torch.tensor([torch.cos(a), -torch.cos(a) * torch.sin(b)]) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == () + res = torch.autograd.functional.jacobian(cost, (a, b)) + assert isinstance(res, tuple) and len(res) == 2 - expected = torch.tensor( - [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)], device=torch_device + expected = ( + torch.tensor([-torch.sin(a), torch.sin(a) * torch.sin(b)]), + torch.tensor([0, -torch.cos(a) * torch.cos(b)]), ) - assert torch.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert torch.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert res[0].shape == (4,) + assert res[1].shape == (4,) - loss = res[0] + res[1] + for _r, _e in zip(res, expected): + assert torch.allclose(_r[:2], _e, atol=atol_for_shots(shots)) + assert torch.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - loss.backward() - expected = torch.tensor( - [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)], - dtype=a.dtype, - device=torch_device, - ) - assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) - assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + else: + assert res[0].shape == (2,) + assert res[1].shape == (2,) + + for _r, _e in zip(res, expected): + assert torch.allclose(_r, _e, atol=atol_for_shots(shots)) - def test_tape_no_parameters(self, torch_device, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1) - params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device) - x, y = params.detach() - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(0.5, wires=0) - qml.expval(qml.PauliZ(0)) + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) + tape4 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, + ) + res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = [qml.math.asarray(ri, like="torch") for r in res for ri in r] + else: + res = [qml.math.asarray(r, like="torch") for r in res] + return sum(torch.hstack(res)) - tape3 = qml.tape.QuantumScript.from_queue(q3) + params = torch.tensor([0.1, 0.2], requires_grad=True) + x, y = params.clone().detach() - res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) - expected = torch.tensor(1 + np.cos(0.5), dtype=res.dtype) + torch.cos(x) * torch.cos(y) - expected = expected.to(device=res.device) + res = cost(params) + expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert torch.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res, 2 * expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) res.backward() - grad = params.grad.detach() - expected = torch.tensor( - [-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)], - dtype=grad.dtype, - device=grad.device, - ) - assert torch.allclose(grad, expected, atol=tol, rtol=0) - - def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = torch.tensor(0.1, requires_grad=True, device=torch_device) - b = torch.tensor(0.2, requires_grad=True, device=torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - assert tape.trainable_params == [0, 1] + expected = torch.tensor([-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)]) + if shots.has_partitioned_shots: + assert torch.allclose(params.grad, 2 * expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(params.grad, expected, atol=atol_for_shots(shots), rtol=0) - res = execute([tape], dev, **execute_kwargs)[0] - loss = res[0] + res[1] - loss.backward() + @pytest.mark.skip("torch cannot reuse tensors in various computations") + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - a_val = 0.54 - b_val = 0.8 - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - b = torch.tensor(b_val, requires_grad=True, device=torch_device) + if execute_kwargs["gradient_fn"] == "backprop": + pytest.xfail("backprop is not compatible with something about this situation.") - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - expected = torch.tensor( - [np.cos(2 * a_val), -np.cos(2 * a_val) * np.sin(b_val)], - device=torch_device, - dtype=res2[0].dtype, - ) - assert torch.allclose(res2[0].detach(), expected[0], atol=tol, rtol=0) - assert torch.allclose(res2[1].detach(), expected[1], atol=tol, rtol=0) + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - loss = res2[0] + res2[1] - loss.backward() + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = execute([tape1, tape2, tape3], device, **execute_kwargs) + return torch.hstack([qml.math.asarray(r, like="torch") for r in res]) - expected = torch.tensor( - [ - -2 * np.sin(2 * a_val) + 2 * np.sin(2 * a_val) * np.sin(b_val), - -np.cos(2 * a_val) * np.cos(b_val), - ], - dtype=a.dtype, - device=torch_device, - ) + params = torch.tensor([0.1, 0.2], requires_grad=True) + x, y = params.clone().detach() - assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) - assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + res = cost(params) + assert isinstance(res, torch.Tensor) + assert res.shape == (4,) - def test_classical_processing(self, torch_device, execute_kwargs, tol): - """Test the classical processing of gate parameters within the quantum tape""" - p_val = [0.1, 0.2] - params = torch.tensor(p_val, requires_grad=True, device=torch_device) + assert torch.allclose(res[0], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) + assert torch.allclose(res[1], torch.tensor(1.0), atol=atol_for_shots(shots)) + assert torch.allclose(res[2], torch.cos(torch.tensor(0.5)), atol=atol_for_shots(shots)) + assert torch.allclose(res[3], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - dev = qml.device("default.qubit.legacy", wires=1) + jac = torch.autograd.functional.jacobian(cost, params) + assert isinstance(jac, torch.Tensor) + assert jac.shape == (4, 2) # pylint: disable=no-member - with qml.queuing.AnnotatedQueue() as q: - qml.RY(params[0] * params[1], wires=0) - qml.RZ(0.2, wires=0) - qml.RX(params[1] + params[1] ** 2 + torch.sin(params[0]), wires=0) - qml.expval(qml.PauliZ(0)) + assert torch.allclose(jac[1:3], torch.tensor(0.0), atol=atol_for_shots(shots)) - tape = qml.tape.QuantumScript.from_queue(q) + d1 = -torch.sin(x) * torch.cos(y) + assert torch.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) # fails for torch + assert torch.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - res = execute([tape], dev, **execute_kwargs)[0] + d2 = -torch.cos(x) * torch.sin(y) + assert torch.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) # fails for torch + assert torch.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - assert tape.trainable_params == [0, 2] + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=True) - tape_params = torch.tensor([i.detach() for i in tape.get_parameters()], device=torch_device) - expected = torch.tensor( - [p_val[0] * p_val[1], p_val[1] + p_val[1] ** 2 + np.sin(p_val[0])], - dtype=tape_params.dtype, - device=torch_device, + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], ) + assert tape.trainable_params == [0, 1] - assert torch.allclose( - tape_params, - expected, - atol=tol, - rtol=0, - ) + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return torch.hstack(execute([new_tape], device, **execute_kwargs)[0]) + + jac = torch.autograd.functional.jacobian(cost, (a, b)) + + a = torch.tensor(0.54, requires_grad=True) + b = torch.tensor(0.8, requires_grad=True) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = torch.tensor([torch.cos(2 * a), -torch.cos(2 * a) * torch.sin(b)]) + assert torch.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) + + jac = torch.autograd.functional.jacobian(lambda a, b: cost(2 * a, b), (a, b)) + expected = ( + torch.tensor([-2 * torch.sin(2 * a), 2 * torch.sin(2 * a) * torch.sin(b)]), + torch.tensor([0, -torch.cos(2 * a) * torch.cos(b)]), + ) + assert isinstance(jac, tuple) and len(jac) == 2 + for _j, _e in zip(jac, expected): + assert torch.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) + + def test_classical_processing(self, execute_kwargs, device, shots): + """Test classical processing within the quantum tape""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=False) + c = torch.tensor(0.3, requires_grad=True) + + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + torch.sin(a), wires=0), + ] - res.backward() + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - assert isinstance(params.grad, torch.Tensor) - assert params.shape == (2,) + # PyTorch docs suggest a lambda for cost functions with some non-trainable args + # See for more: https://pytorch.org/docs/stable/autograd.html#functional-higher-level-api + res = torch.autograd.functional.jacobian(lambda _a, _c: cost(_a, b, _c), (a, c)) - def test_no_trainable_parameters(self, torch_device, execute_kwargs): - """Test evaluation and Jacobian if there are no trainable parameters""" - dev = qml.device("default.qubit.legacy", wires=2) + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - with qml.queuing.AnnotatedQueue() as q: - qml.RY(0.2, wires=0) - qml.RX(torch.tensor(0.1, device=torch_device), wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - tape = qml.tape.QuantumScript.from_queue(q) + @pytest.mark.skip("torch handles gradients and jacobians differently") + def test_no_trainable_parameters(self, execute_kwargs, shots, device): + """Test evaluation and Jacobian if there are no trainable parameters""" + a = torch.tensor(0.1, requires_grad=False) + b = torch.tensor(0.2, requires_grad=False) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [] + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - assert isinstance(res, tuple) - assert len(res) == 2 + res = cost(a, b) + assert res.shape == (2,) - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == () + res = torch.autograd.functional.jacobian(cost, (a, b)) + assert len(res) == 0 - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == () + def loss(a, b): + return torch.sum(cost(a, b)) - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res[0].backward() + res = loss(a, b) + res.backward() - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res[1].backward() + assert torch.allclose(torch.tensor([a.grad, b.grad]), 0) - @pytest.mark.parametrize( - "U", [torch.tensor([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]])] - ) - def test_matrix_parameter(self, torch_device, U, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the torch interface works correctly with a matrix parameter""" - a_val = 0.1 - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - - if isinstance(U, torch.Tensor) and torch_device is not None: - U = U.to(torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + U = torch.tensor([[0, 1], [1, 0]], requires_grad=False, dtype=torch.float64) + a = torch.tensor(0.1, requires_grad=True) - tape = qml.tape.QuantumScript.from_queue(q) + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1] + res = cost(a, U) + assert torch.allclose(res, -torch.cos(a), atol=atol_for_shots(shots)) - expected = torch.tensor(-np.cos(a_val), dtype=res.dtype, device=torch_device) - assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + jac = torch.autograd.functional.jacobian(lambda y: cost(y, U), a) + assert isinstance(jac, torch.Tensor) + assert torch.allclose(jac, torch.sin(a), atol=atol_for_shots(shots), rtol=0) - res.backward() - expected = torch.tensor([np.sin(a_val)], dtype=a.grad.dtype, device=torch_device) - assert torch.allclose(a.grad, expected, atol=tol, rtol=0) - - def test_differentiable_expand(self, torch_device, execute_kwargs, tol): - """Test that operation and nested tape expansion + def test_differentiable_expand(self, execute_kwargs, device, shots): + """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -716,519 +507,244 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device("default.qubit.legacy", wires=1) - a = np.array(0.1) - p_val = [0.1, 0.2, 0.3] - p = torch.tensor(p_val, requires_grad=True, device=torch_device) - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - res = execute([tape], dev, **execute_kwargs)[0] + a = torch.tensor(0.1, requires_grad=False) + p = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - expected = torch.tensor( - np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) - + np.sin(a) - * ( - np.cos(p_val[2]) * np.sin(p_val[1]) - + np.cos(p_val[0]) * np.cos(p_val[1]) * np.sin(p_val[2]) - ), - dtype=res.dtype, - device=torch_device, + res = cost_fn(a, p) + expected = torch.cos(a) * torch.cos(p[1]) * torch.sin(p[0]) + torch.sin(a) * ( + torch.cos(p[2]) * torch.sin(p[1]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.sin(p[2]) ) - assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res.backward() + res = torch.autograd.functional.jacobian(lambda _p: cost_fn(a, _p), p) expected = torch.tensor( [ - np.cos(p_val[1]) - * (np.cos(a) * np.cos(p_val[0]) - np.sin(a) * np.sin(p_val[0]) * np.sin(p_val[2])), - np.cos(p_val[1]) * np.cos(p_val[2]) * np.sin(a) - - np.sin(p_val[1]) - * (np.cos(a) * np.sin(p_val[0]) + np.cos(p_val[0]) * np.sin(a) * np.sin(p_val[2])), - np.sin(a) + torch.cos(p[1]) + * ( + torch.cos(a) * torch.cos(p[0]) + - torch.sin(a) * torch.sin(p[0]) * torch.sin(p[2]) + ), + torch.cos(p[1]) * torch.cos(p[2]) * torch.sin(a) + - torch.sin(p[1]) + * ( + torch.cos(a) * torch.sin(p[0]) + + torch.cos(p[0]) * torch.sin(a) * torch.sin(p[2]) + ), + torch.sin(a) * ( - np.cos(p_val[0]) * np.cos(p_val[1]) * np.cos(p_val[2]) - - np.sin(p_val[1]) * np.sin(p_val[2]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.cos(p[2]) + - torch.sin(p[1]) * torch.sin(p[2]) ), - ], - dtype=p.grad.dtype, - device=torch_device, + ] ) - assert torch.allclose(p.grad, expected, atol=tol, rtol=0) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, torch_device, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - - dev = qml.device("default.qubit.legacy", wires=2) - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True, device=torch_device) - y = torch.tensor(y_val, requires_grad=True, device=torch_device) - - def circuit(x, y): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] - - res = circuit(x, y) - - expected_0 = torch.tensor( - [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2], - dtype=res[0].dtype, - device=torch_device, - ) + x = torch.tensor(0.543, requires_grad=True) + y = torch.tensor(-0.654, requires_grad=True) - expected_1 = torch.tensor( + res = cost(x, y) + expected = torch.tensor( [ - (1 + np.cos(x_val) * np.cos(y_val)) / 2, - (1 - np.cos(x_val) * np.cos(y_val)) / 2, - ], - dtype=res[0].dtype, - device=torch_device, + [ + torch.cos(x / 2) ** 2, + torch.sin(x / 2) ** 2, + (1 + torch.cos(x) * torch.cos(y)) / 2, + (1 - torch.cos(x) * torch.cos(y)) / 2, + ], + ] ) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[0], expected_0, atol=tol, rtol=0) - assert torch.allclose(res[1], expected_1, atol=tol, rtol=0) - - jac = torch_functional.jacobian(circuit, (x, y)) - dtype_jac = jac[0][0].dtype + res = torch.autograd.functional.jacobian(cost, (x, y)) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) - res_0 = torch.tensor( - [-np.sin(x_val) / 2, np.sin(x_val) / 2], dtype=dtype_jac, device=torch_device - ) - res_1 = torch.tensor([0.0, 0.0], dtype=dtype_jac, device=torch_device) - res_2 = torch.tensor( - [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], - dtype=dtype_jac, - device=torch_device, - ) - res_3 = torch.tensor( - [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], - dtype=dtype_jac, - device=torch_device, + expected = ( + torch.tensor( + [ + [ + -torch.sin(x) / 2, + torch.sin(x) / 2, + -torch.sin(x) * torch.cos(y) / 2, + torch.sin(x) * torch.cos(y) / 2, + ], + ] + ), + torch.tensor( + [ + [0, 0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2], + ] + ), ) - assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) - assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) - assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) - assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - def test_ragged_differentiation(self, torch_device, execute_kwargs, tol): + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - dev = qml.device("default.qubit.legacy", wires=2) - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True, device=torch_device) - y = torch.tensor(y_val, requires_grad=True, device=torch_device) + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - def circuit(x, y): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) + x = torch.tensor(0.543, requires_grad=True) + y = torch.tensor(-0.654, requires_grad=True) - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] - - res = circuit(x, y) - - res_0 = torch.tensor(np.cos(x_val), dtype=res[0].dtype, device=torch_device) - res_1 = torch.tensor( - [(1 + np.cos(x_val) * np.cos(y_val)) / 2, (1 - np.cos(x_val) * np.cos(y_val)) / 2], - dtype=res[0].dtype, - device=torch_device, + res = cost(x, y) + expected = torch.tensor( + [ + torch.cos(x), + (1 + torch.cos(x) * torch.cos(y)) / 2, + (1 - torch.cos(x) * torch.cos(y)) / 2, + ] ) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert torch.allclose(res[0], res_0, atol=tol, rtol=0) - assert torch.allclose(res[1], res_1, atol=tol, rtol=0) - - jac = torch_functional.jacobian(circuit, (x, y)) - dtype_jac = jac[0][0].dtype + res = torch.autograd.functional.jacobian(cost, (x, y)) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) - res_0 = torch.tensor( - -np.sin(x_val), - dtype=dtype_jac, - device=torch_device, - ) - res_1 = torch.tensor( - 0.0, - dtype=dtype_jac, - device=torch_device, - ) - res_2 = torch.tensor( - [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], - dtype=dtype_jac, - device=torch_device, - ) - res_3 = torch.tensor( - [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], - dtype=dtype_jac, - device=torch_device, + expected = ( + torch.tensor( + [-torch.sin(x), -torch.sin(x) * torch.cos(y) / 2, torch.sin(x) * torch.cos(y) / 2] + ), + torch.tensor([0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2]), ) - - assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) - assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) - assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) - assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) - - def test_sampling(self, torch_device, execute_kwargs): - """Test sampling works as expected""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - if execute_kwargs["interface"] == "auto": - pytest.skip("Can't detect interface without a parametrized gate in the tape") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (10,) - - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == (10,) - - def test_sampling_expval(self, torch_device, execute_kwargs): - """Test sampling works as expected if combined with expectation values""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - if execute_kwargs["interface"] == "auto": - pytest.skip("Can't detect interface without a parametrized gate in the tape") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.expval(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - assert len(res) == 2 - assert isinstance(res, tuple) - assert res[0].shape == (10,) - assert res[1].shape == () - assert isinstance(res[0], torch.Tensor) - assert isinstance(res[1], torch.Tensor) - - def test_sampling_gradient_error(self, torch_device, execute_kwargs): - """Test differentiating a tape with sampling results in an error""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - - dev = qml.device("default.qubit.legacy", wires=1, shots=10) - - x = torch.tensor(0.65, requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.sample() - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res.backward() - - def test_repeated_application_after_expand(self, torch_device, execute_kwargs): - """Test that the Torch interface continues to work after - tape expansions""" - # pylint: disable=unused-argument - n_qubits = 2 - dev = qml.device("default.qubit.legacy", wires=n_qubits) - - weights = torch.ones((3,)) - - with qml.queuing.AnnotatedQueue() as q: - qml.U3(*weights, wires=0) - qml.expval(qml.PauliZ(wires=0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - tape = tape.expand() - execute([tape], dev, **execute_kwargs) + assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) -@pytest.mark.parametrize("torch_device", torch_devices) class TestHigherOrderDerivatives: """Test that the torch execute function can be differentiated""" @pytest.mark.parametrize( "params", [ - torch.tensor([0.543, -0.654], requires_grad=True), - torch.tensor([0, -0.654], requires_grad=True), - torch.tensor([-2.0, 0], requires_grad=True), + torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64), + torch.tensor([0, -0.654], requires_grad=True, dtype=torch.float64), + torch.tensor([-2.0, 0], requires_grad=True, dtype=torch.float64), ], ) - def test_parameter_shift_hessian(self, torch_device, params, tol): + def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using torch, yielding second derivatives.""" - # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, interface="torch", max_diff=2 - ) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] res = cost_fn(params) - x, y = params.detach() - expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + x, y = params.clone().detach() + expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) assert torch.allclose(res, expected, atol=tol, rtol=0) res.backward() expected = torch.tensor( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] ) - assert torch.allclose(params.grad.detach(), expected, atol=tol, rtol=0) + assert torch.allclose(params.grad, expected, atol=tol, rtol=0) res = torch.autograd.functional.hessian(cost_fn, params) expected = torch.tensor( [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [-torch.cos(2 * x) * torch.cos(2 * y), torch.sin(2 * x) * torch.sin(2 * y)], + [torch.sin(2 * x) * torch.sin(2 * y), -2 * torch.cos(x) ** 2 * torch.cos(2 * y)], ] ) assert torch.allclose(res, expected, atol=tol, rtol=0) - def test_hessian_vector_valued(self, torch_device, tol): - """Test hessian calculation of a vector valued tape""" - dev = qml.device("default.qubit.legacy", wires=1) - - def circuit(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return torch.stack( - execute([tape], dev, gradient_fn=param_shift, interface="torch", max_diff=2) - ) - - x = torch.tensor([1.0, 2.0], requires_grad=True, device=torch_device) - res = circuit(x) - - if torch_device is not None: - a, b = x.detach().cpu().numpy() - else: - a, b = x.detach().numpy() - - expected_res = torch.tensor( - [ - 0.5 + 0.5 * np.cos(a) * np.cos(b), - 0.5 - 0.5 * np.cos(a) * np.cos(b), - ], - dtype=res.dtype, - device=torch_device, - ) - assert torch.allclose(res.detach(), expected_res, atol=tol, rtol=0) - - def jac_fn(x): - return torch_functional.jacobian(circuit, x, create_graph=True) - - g = jac_fn(x) - - hess = torch_functional.jacobian(jac_fn, x) - - expected_g = torch.tensor( - [ - [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)], - [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)], - ], - dtype=g.dtype, - device=torch_device, - ) - assert torch.allclose(g.detach(), expected_g, atol=tol, rtol=0) - - expected_hess = torch.tensor( - [ - [ - [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)], - [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)], - ], - [ - [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)], - [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)], - ], - ], - dtype=hess.dtype, - device=torch_device, - ) - assert torch.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - - def test_adjoint_hessian(self, torch_device, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor( - [0.543, -0.654], requires_grad=True, dtype=torch.float64, device=torch_device - ) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface="torch", - )[0] - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2], dtype=torch.float64, device=torch_device) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, torch_device, tol): + def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + dev = DefaultQubit() + params = torch.tensor([0.543, -0.654], requires_grad=True) def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface="torch" - ) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) return result[0] + result[1][0] res = cost_fn(params) - x, y = params.detach() - expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) - assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + x, y = params.clone().detach() + expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) + assert torch.allclose(res, expected, atol=tol, rtol=0) res.backward() expected = torch.tensor( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert torch.allclose( - params.grad.detach().to(torch_device), expected.to(torch_device), atol=tol, rtol=0 + [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] ) + assert torch.allclose(params.grad, expected, atol=tol, rtol=0) res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2], dtype=torch.float64) - assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) - - -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift, "interface": "torch"}, - {"gradient_fn": finite_diff, "interface": "torch"}, -] + expected = torch.zeros([2, 2]) + assert torch.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -1237,72 +753,97 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - - return torch.hstack(execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + res = execute([tape], device, **execute_kwargs)[0] + if shots.has_partitioned_shots: + return torch.hstack(res[0] + res[1]) + return torch.hstack(res) return _cost_fn @staticmethod def cost_fn_expected(weights, coeffs1, coeffs2): """Analytic value of cost_fn above""" - a, b, c = coeffs1.detach().numpy() - d = coeffs2.detach().numpy()[0] - x, y = weights.detach().numpy() - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + a, b, c = coeffs1.clone().detach() + d = coeffs2[0].clone().detach() + x, y = weights.clone().detach() + return torch.tensor( + [ + -c * torch.sin(x) * torch.sin(y) + torch.cos(x) * (a + b * torch.sin(y)), + d * torch.cos(x), + ] + ) @staticmethod def cost_fn_jacobian(weights, coeffs1, coeffs2): """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.detach().numpy() - d = coeffs2.detach().numpy()[0] - x, y = weights.detach().numpy() - return np.array( + a, b, c = coeffs1.clone().detach() + d = coeffs2[0].clone().detach() + x, y = weights.clone().detach() + return torch.tensor( [ [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), + -c * torch.cos(x) * torch.sin(y) - torch.sin(x) * (a + b * torch.sin(y)), + b * torch.cos(x) * torch.cos(y) - c * torch.cos(y) * torch.sin(x), + torch.cos(x), + torch.cos(x) * torch.sin(y), + -(torch.sin(x) * torch.sin(y)), 0, ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + [-d * torch.sin(x), 0, 0, 0, 0, torch.cos(x)], ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False, dtype=torch.float64) - coeffs2 = torch.tensor([0.7], requires_grad=False, dtype=torch.float64) - weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) - dev = qml.device("default.qubit.legacy", wires=2) + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False) + coeffs2 = torch.tensor([0.7], requires_grad=False) + weights = torch.tensor([0.4, 0.5], requires_grad=True) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = torch.hstack( - torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) - ) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + res = torch.autograd.functional.jacobian(lambda w: cost_fn(w, coeffs1, coeffs2), weights) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + if shots.has_partitioned_shots: + assert torch.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64) - coeffs2 = torch.tensor([0.7], requires_grad=True, dtype=torch.float64) - weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) - dev = qml.device("default.qubit.legacy", wires=2) + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) + coeffs2 = torch.tensor([0.7], requires_grad=True) + weights = torch.tensor([0.4, 0.5], requires_grad=True) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = torch.hstack( - torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) - ) + res = torch.hstack(torch.autograd.functional.jacobian(cost_fn, (weights, coeffs1, coeffs2))) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_torch_qnode.py b/tests/interfaces/test_torch_qnode.py index 3965255a4a4..2f1c85f6275 100644 --- a/tests/interfaces/test_torch_qnode.py +++ b/tests/interfaces/test_torch_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 Xanadu Quantum Technologies Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the Torch interface with a QNode""" -# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods, -# pylint: disable=use-dict-literal, use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment +# pylint: disable=too-many-arguments,unexpected-keyword-arg,no-member,comparison-with-callable, no-name-in-module +# pylint: disable=use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment, use-dict-literal import numpy as np import pytest import pennylane as qml from pennylane import qnode +from pennylane.devices import DefaultQubit +from tests.param_shift_dev import ParamShiftDerivativesDevice pytestmark = pytest.mark.torch @@ -26,14 +28,25 @@ jacobian = torch.autograd.functional.jacobian hessian = torch.autograd.functional.hessian +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", True, True], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, False], + [ParamShiftDerivativesDevice(), "best", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, True], ] interface_and_qubit_device_and_diff_method = [ @@ -46,25 +59,25 @@ @pytest.mark.parametrize( - "interface, dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface, dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with Torch integrates with the PennyLane stack""" - def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): + def test_execution_with_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -79,7 +92,7 @@ def circuit(a): # with the interface, the tape returns torch tensors assert isinstance(res, torch.Tensor) - assert res.shape == tuple() + assert res.shape == () # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -89,20 +102,18 @@ def circuit(a): grad = a.grad assert isinstance(grad, torch.Tensor) - assert grad.shape == tuple() + assert grad.shape == () - def test_interface_swap(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the Torch interface can be applied to a QNode with a pre-existing interface""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, interface="autograd", grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface="autograd", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -128,22 +139,19 @@ def circuit(a): assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, interface, dev_name, diff_method, grad_on_execution): + def test_drawing(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test circuit drawing when using the torch interface""" x = torch.tensor(0.1, requires_grad=True) y = torch.tensor([0.2, 0.3], requires_grad=True) z = torch.tensor(0.4, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) @@ -159,14 +167,17 @@ def circuit(p1, p2=y, **kwargs): assert result == expected - def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test jacobian calculation""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) - if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a_val = 0.1 @@ -175,13 +186,6 @@ def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -217,19 +221,29 @@ def circuit(a, b): assert np.allclose(b.grad, expected[1], atol=tol, rtol=0) # TODO: fix this behavior with float: already present before return type. - @pytest.mark.xfail - def test_jacobian_dtype(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_dtype( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test calculating the jacobian with a different datatype""" + if not "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Failing unless lightning.qubit") if diff_method == "backprop": pytest.skip("Test does not support backprop") a = torch.tensor(0.1, dtype=torch.float32, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True) - dev = qml.device(dev_name, wires=2) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -253,15 +267,20 @@ def circuit(a, b): assert a.grad.dtype is torch.float32 assert b.grad.dtype is torch.float32 - def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_options( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test setting jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite-diff and spsa") a = torch.tensor([0.1, 0.2], requires_grad=True) - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, @@ -269,6 +288,7 @@ def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execut interface=interface, h=1e-8, approx_order=2, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -278,7 +298,9 @@ def circuit(a): res = circuit(a) res.backward() - def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_changing_trainability( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -290,10 +312,12 @@ def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_e a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) - dev = qml.device(dev_name, wires=2) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -342,21 +366,25 @@ def circuit(a, b): expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val) assert np.allclose(a.grad, expected, atol=tol, rtol=0) - def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): + def test_classical_processing( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test classical processing within the quantum tape""" a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False) c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -376,17 +404,22 @@ def circuit(a, b, c): assert b.grad is None assert isinstance(c.grad, torch.Tensor) - def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): + def test_no_trainable_parameters( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test evaluation and Jacobian if there are no trainable parameters""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -430,21 +463,20 @@ def circuit(a, b): np.array([[0, 1], [1, 0]]), ], ) - def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, U, tol): + def test_matrix_parameter( + self, interface, dev, diff_method, grad_on_execution, device_vjp, U, tol + ): """Test that the Torch interface works correctly with a matrix parameter""" a_val = 0.1 a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -461,18 +493,23 @@ def circuit(U, a): res.backward() assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0) - def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_differentiable_expand( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that operation and nested tapes expansion is differentiable""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - class U3(qml.U3): + class U3(qml.U3): # pylint:disable=too-few-public-methods def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -481,12 +518,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) a = np.array(0.1) p_val = [0.1, 0.2, 0.3] p = torch.tensor(p_val, dtype=torch.float64, requires_grad=True) @@ -527,9 +558,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, mocker, tol): + def test_changing_shots(self): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -537,33 +568,22 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # only same call as above + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # TODO: add this test after shot vectors addition @pytest.mark.xfail def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - # pylint: disable=unexpected-keyword-arg - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = torch.tensor([0.543, -0.654], requires_grad=True) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -585,11 +605,10 @@ def test_multiple_gradient_integration(self, tol): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) weights = torch.tensor([0.543, -0.654], requires_grad=True) a, b = weights - @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) + @qnode(DefaultQubit(), interface="torch", diff_method=qml.gradients.param_shift) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -597,7 +616,7 @@ def circuit(a, b): return qml.expval(qml.PauliY(1)) res1 = circuit(*weights) - assert qml.math.shape(res1) == tuple() + assert qml.math.shape(res1) == () res2 = circuit(*weights, shots=[(1, 1000)]) assert len(res2) == 1000 @@ -609,120 +628,47 @@ def circuit(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = torch.tensor([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(dev, interface="torch") + @qnode(DefaultQubit(), interface="torch") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - - # original QNode settings are unaffected - + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - - -class TestAdjoint: - """Specific integration tests for the adjoint method""" - - def test_reuse_state(self, mocker): - """Tests that the Torch interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, diff_method="adjoint", interface="torch") - def circ(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=1) - qml.CNOT(wires=(0, 1)) - return qml.expval(qml.PauliZ(0)) - - expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])]) - - spy = mocker.spy(dev, "adjoint_jacobian") - - x1 = torch.tensor([1.0, 1.0], requires_grad=True) - res = circ(x1) - res.backward() - - assert np.allclose(x1.grad[0], expected_grad(x1)[0]) - assert circ.device.num_executions == 1 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) - - def test_resuse_state_multiple_evals(self, mocker, tol): - """Tests that the Torch interface reuses the device state for adjoint differentiation, - even where there are intermediate evaluations.""" - dev = qml.device("default.qubit.legacy", wires=2) - - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True) - y = torch.tensor(y_val, requires_grad=True) - - @qnode(dev, diff_method="adjoint", interface="torch") - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - res1 = circuit(x, y) - assert np.allclose(res1.detach(), np.cos(x_val), atol=tol, rtol=0) - - # intermediate evaluation with different values - circuit(torch.tan(x), torch.cosh(y)) - - # the adjoint method will continue to compute the correct derivative - res1.backward() - assert np.allclose(x.grad.detach(), -np.sin(x_val), atol=tol, rtol=0) - assert dev.num_executions == 2 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = {} - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -733,6 +679,7 @@ def test_probability_differentiation( diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -768,25 +715,24 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_ragged_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -824,17 +770,22 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit1(weights): qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) @@ -866,18 +817,11 @@ def cost(weights): loss = cost(weights) loss.backward() - def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a scalar valued QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -888,6 +832,7 @@ def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol) grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -923,18 +868,13 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_vector_valued( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test hessian calculation of a vector valued QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -945,6 +885,7 @@ def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_e grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -991,18 +932,11 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a ragged QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -1013,6 +947,7 @@ def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_executio grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -1071,29 +1006,23 @@ def circuit_stack(x): assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") options = {} if diff_method == "finite-diff": options = {"h": 1e-6} - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -1138,23 +1067,23 @@ def cost_fn(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.skip("parameter shift does not support measuring the state.") + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning devices do not support state with adjoint diff.") x = torch.tensor(0.543, requires_grad=True) y = torch.tensor(-0.654, requires_grad=True) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1164,7 +1093,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is torch.complex128 + assert torch.is_complex(res) probs = torch.abs(res) ** 2 return probs[0] + probs[2] @@ -1181,20 +1110,25 @@ def cost_fn(x, y): assert torch.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): + def test_projector( + self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the variance of a projector is correctly returned""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "spsa": + pytest.skip("adjoint supports either all expvals or all diagonal measurements") + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard does not support variances.") - dev = qml.device(dev_name, wires=2) P = torch.tensor(state, requires_grad=False) x, y = 0.765, -0.654 @@ -1222,99 +1156,64 @@ def circuit(x, y): ) assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = torch.tensor(0.543, dtype=torch.float64, requires_grad=True) - phi = torch.tensor(-0.654, dtype=torch.float64, requires_grad=True) - - @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = torch.exp(2 * r) * torch.sin(phi) ** 2 + torch.exp(-2 * r) * torch.cos(phi) ** 2 - assert torch.allclose(res, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - # circuit jacobians - res.backward() - res = torch.tensor([r.grad, phi.grad]) - expected = torch.tensor( - [ - [ - 2 * torch.exp(2 * r) * torch.sin(phi) ** 2 - - 2 * torch.exp(-2 * r) * torch.cos(phi) ** 2, - 2 * torch.sinh(2 * r) * torch.sin(2 * phi), - ] - ] + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - n = torch.tensor(0.12, dtype=torch.float64, requires_grad=True) - a = torch.tensor(0.765, dtype=torch.float64, requires_grad=True) + phi = torch.tensor(1.23, requires_grad=True) + theta = torch.tensor(4.56, requires_grad=True) - @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) + assert qml.math.allclose(circuit(phi, theta), expected_circuit(theta)) - res = circuit(n, a) - expected = n**2 + n + torch.abs(a) ** 2 * (1 + 2 * n) - assert torch.allclose(res, expected, atol=tol, rtol=0) + gradient = torch.autograd.grad(circuit(phi, theta), [phi, theta]) + exp_theta_grad = torch.autograd.grad(expected_circuit(theta), theta)[0] + assert qml.math.allclose(gradient, [0.0, exp_theta_grad]) - # circuit jacobians - res.backward() - res = torch.tensor([n.grad, a.grad]) - expected = torch.tensor([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) - assert torch.allclose(res, expected, atol=tol, rtol=0) - -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method +) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Torch interface""" - def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution): + def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1325,6 +1224,7 @@ def decomposition(self): diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, interface="torch", ) def circuit(x): @@ -1335,33 +1235,32 @@ def circuit(x): x = torch.tensor(0.5, requires_grad=True, dtype=torch.float64) loss = circuit(x) + loss.backward() res = x.grad assert torch.allclose(res, -3 * torch.sin(3 * x)) - if diff_method == "parameter-shift": + if diff_method == "parameter-shift" and dev.name != "param_shift.qubit": # test second order derivatives res = torch.autograd.functional.hessian(circuit, x) assert torch.allclose(res, -9 * torch.cos(3 * x)) @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff + self, + dev, + diff_method, + grad_on_execution, + device_vjp, + max_diff, ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1373,6 +1272,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1388,7 +1288,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol + self, dev, diff_method, grad_on_execution, max_diff, device_vjp, tol ): """Test that if there are non-commuting groups and the number of shots is None @@ -1398,17 +1298,17 @@ def test_hamiltonian_expansion_analytic( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, ) if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1448,7 +1348,11 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method in ("parameter-shift", "backprop") and max_diff == 2: + if ( + diff_method in ("parameter-shift", "backprop") + and max_diff == 2 + and dev.name != "param_shift.qubit" + ): hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] @@ -1468,14 +1372,14 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, device_vjp, grad_on_execution, max_diff ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1487,11 +1391,10 @@ def test_hamiltonian_expansion_finite_shots( tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + tol = 0.15 elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1500,6 +1403,7 @@ def test_hamiltonian_expansion_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1514,15 +1418,17 @@ def circuit(data, weights, coeffs): c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin( d[1] + w[1] ) assert torch.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + res.backward() grad = (w.grad, c.grad) @@ -1541,8 +1447,10 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) + if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": + hessians = torch.autograd.functional.hessian( + lambda _d, _w, _c: circuit(_d, _w, _c, shots=50000), (d, w, c) + ) grad2_c = hessians[2][2] assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64), atol=tol) @@ -1566,15 +1474,14 @@ class TestSample: def test_sample_dimension(self): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert len(res) == 2 @@ -1587,37 +1494,33 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (shots,) + assert res[0].shape == (10,) assert isinstance(res[1], torch.Tensor) assert res[1].shape == () def test_counts_expval(self): """Test counts works as expected if combined with expectation values""" - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(qml.device("default.qubit"), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) @@ -1628,22 +1531,19 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, tuple) assert len(result) == 3 - assert np.array_equal(result[0].shape, (n_sample,)) + assert np.array_equal(result[0].shape, (10,)) assert result[1].shape == () assert isinstance(result[1], torch.Tensor) assert result[2].shape == () @@ -1652,73 +1552,73 @@ def circuit(): def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" - n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, torch.Tensor) - assert np.array_equal(result.shape, (n_sample,)) + assert np.array_equal(result.shape, (10,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=10) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) - assert tuple(result[0].shape) == (n_sample,) - assert tuple(result[1].shape) == (n_sample,) - assert tuple(result[2].shape) == (n_sample,) + assert tuple(result[0].shape) == (10,) + assert tuple(result[1].shape) == (10,) + assert tuple(result[2].shape) == (10,) assert result[0].dtype == torch.float64 assert result[1].dtype == torch.float64 assert result[2].dtype == torch.float64 qubit_device_and_diff_method_and_grad_on_execution = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", True, True], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "hadamard", False, False], ] @pytest.mark.parametrize( - "dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method_and_grad_on_execution + "dev,diff_method,grad_on_execution, device_vjp", + qubit_device_and_diff_method_and_grad_on_execution, ) @pytest.mark.parametrize("shots", [None, 10000]) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution, shots): + # pylint:disable=too-many-public-methods + + def test_grad_single_measurement_param( + self, dev, diff_method, grad_on_execution, device_vjp, shots + ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1726,7 +1626,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - res = circuit(a) + res = circuit(a, shots=shots) assert isinstance(res, torch.Tensor) assert res.shape == () @@ -1738,20 +1638,19 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, shots, device_vjp ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1760,7 +1659,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - res = circuit(a, b) + res = circuit(a, b, shots=shots) # gradient res.backward() @@ -1771,20 +1670,19 @@ def circuit(a, b): assert grad_b.shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1792,7 +1690,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - res = circuit(a) + res = circuit(a, shots=shots) # gradient res.backward() @@ -1802,7 +1700,7 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1812,14 +1710,13 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1833,7 +1730,7 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1842,14 +1739,13 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1858,7 +1754,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(circuit, (a, b)) + jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) assert isinstance(jac, tuple) @@ -1869,7 +1765,7 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1878,39 +1774,31 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, torch.Tensor) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1919,6 +1807,7 @@ def test_jacobian_expval_expval_multiple_params( interface="torch", diff_method=diff_method, max_diff=1, + device_vjp=device_vjp, grad_on_execution=grad_on_execution, ) def circuit(x, y): @@ -1927,7 +1816,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, (par_0, par_1)) + jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(jac, tuple) @@ -1946,20 +1835,19 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1967,7 +1855,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1979,7 +1867,7 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, device_vjp, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1989,8 +1877,6 @@ def test_jacobian_var_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1998,6 +1884,7 @@ def test_jacobian_var_var_multiple_params( dev, interface="torch", diff_method=diff_method, + device_vjp=device_vjp, max_diff=1, grad_on_execution=grad_on_execution, ) @@ -2007,7 +1894,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, (par_0, par_1)) + jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2027,7 +1914,7 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, device_vjp, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2037,9 +1924,13 @@ def test_jacobian_var_var_multiple_params_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2047,7 +1938,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2059,22 +1950,22 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -2082,7 +1973,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2094,7 +1985,7 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2102,14 +1993,13 @@ def test_jacobian_multiple_measurement_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2118,7 +2008,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(circuit, (a, b)) + jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2138,7 +2028,7 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2146,14 +2036,13 @@ def test_jacobian_multiple_measurement_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2161,7 +2050,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2172,23 +2061,18 @@ def circuit(a): assert isinstance(jac[1], torch.Tensor) assert jac[1].shape == (4, 2) - def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): + def test_hessian_expval_multiple_params( + self, dev, diff_method, grad_on_execution, shots, device_vjp + ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2196,6 +2080,7 @@ def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_exe diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2203,7 +2088,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, (par_0, par_1)) + hess = hessian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2223,7 +2108,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2231,13 +2116,6 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( @@ -2246,6 +2124,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2253,12 +2132,14 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, params) + hess = hessian(lambda _a: circuit(_a, shots=shots), params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): + def test_hessian_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp, shots + ): """The hessian of a single measurement with multiple params returns a tuple of arrays.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2267,15 +2148,8 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2283,6 +2157,7 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2290,7 +2165,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, (par_0, par_1)) + hess = hessian(lambda _a, _b: circuit(_a, _b, shots=shots), (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2310,7 +2185,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2320,9 +2195,7 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - - params = torch.tensor([0.1, 0.2], requires_grad=True) + params = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2330,6 +2203,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2337,13 +2211,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, params) + hess = hessian(lambda _a: circuit(_a, shots=shots), params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2354,15 +2228,8 @@ def test_hessian_probs_expval_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2370,6 +2237,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2378,7 +2246,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y)) + return torch.hstack(circuit(x, y, shots=shots)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2402,10 +2270,9 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2414,7 +2281,7 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par = torch.tensor([0.1, 0.2], requires_grad=True) + par = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2422,6 +2289,7 @@ def test_hessian_expval_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2430,7 +2298,7 @@ def circuit(x): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x)) + return torch.hstack(circuit(x, shots=shots)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2440,10 +2308,9 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2461,6 +2328,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2469,7 +2337,7 @@ def circuit(x, y): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y)) + return torch.hstack(circuit(x, y, shots=shots)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2493,10 +2361,9 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_var_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2513,6 +2380,7 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2521,7 +2389,7 @@ def circuit(x): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x)) + return torch.hstack(circuit(x, shots=shots)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2531,14 +2399,11 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="torch") + @qml.qnode(DefaultQubit(), interface="torch") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_transform_program_integration.py b/tests/interfaces/test_transform_program_integration.py index 6d254f953dc..ef492b52035 100644 --- a/tests/interfaces/test_transform_program_integration.py +++ b/tests/interfaces/test_transform_program_integration.py @@ -17,18 +17,20 @@ """ import copy from functools import partial -from typing import Callable, Tuple import numpy as np import pytest import pennylane as qml +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn -device_suite = ( - qml.device("default.qubit.legacy", wires=5), - qml.devices.DefaultQubit(), - qml.device("lightning.qubit", wires=5), -) +with pytest.warns(qml.PennyLaneDeprecationWarning): + device_suite = ( + qml.device("default.qubit.legacy", wires=5), + qml.devices.DefaultQubit(), + qml.device("lightning.qubit", wires=5), + ) @pytest.mark.all_interfaces @@ -66,7 +68,7 @@ def null_postprocessing(results): def just_pauli_x_out( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing @@ -168,19 +170,21 @@ def split_sum_terms(tape): def test_chained_preprocessing(self): """Test a transform program with two transforms where their order affects the output.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) def null_postprocessing(results): return results[0] - def just_pauli_x_out(tape: qml.tape.QuantumTape) -> (Tuple[qml.tape.QuantumTape], Callable): + def just_pauli_x_out( + tape: qml.tape.QuantumTape, + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing def repeat_operations( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: new_tape = qml.tape.QuantumScript( tape.operations + copy.deepcopy(tape.operations), tape.measurements ) diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py index 5b57bca85a4..b90c59f5072 100644 --- a/tests/logging/test_logging_autograd.py +++ b/tests/logging/test_logging_autograd.py @@ -123,7 +123,7 @@ def circuit(params): [ "Creating QNode(func=.""" + basis_state = [0] + data = [1] + row_indices = [0] + col_indices = [0] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(2, 2)) + + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_single_qubit_basis_state_1(self): + """Tests the function with a single-qubit basis state |1>.""" + basis_state = [1] + data = [1] + row_indices = [1] + col_indices = [1] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(2, 2)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_10(self): + """Tests the function with a two-qubits basis state |10>.""" + basis_state = [1, 0] + data = [1] + row_indices = [2] + col_indices = [2] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_01(self): + """Tests the function with a two-qubits basis state |01>.""" + basis_state = [0, 1] + data = [1] + row_indices = [1] + col_indices = [1] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_11(self): + """Tests the function with a two-qubits basis state |11>.""" + basis_state = [1, 1] + data = [1] + row_indices = [3] + col_indices = [3] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_three_qubit_basis_state_101(self): + """Tests the function with a three-qubits basis state |101>.""" + basis_state = [1, 0, 1] + data = [1] + row_indices = [5] + col_indices = [5] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(8, 8)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_invalid_basis_state(self): + """Tests the function with an invalid state.""" + basis_state = [0, 2] # Invalid basis state + with pytest.raises(ValueError): + BasisStateProjector.compute_sparse_matrix(basis_state) + @pytest.mark.jax def test_jit_measurement(self): """Test that the measurement of a projector can be jitted.""" @@ -676,10 +746,10 @@ def test_matrix_representation(self, basis_state, expected, n_wires, tol): assert np.allclose(res_dynamic, expected, atol=tol) assert np.allclose(res_static, expected, atol=tol) - @pytest.mark.parametrize( - "dev", (qml.device("default.qubit"), qml.device("default.qubit.legacy", wires=1)) - ) - def test_integration_batched_state(self, dev): + @pytest.mark.parametrize("dev_name", ("default.qubit", "default.qubit.legacy")) + def test_integration_batched_state(self, dev_name): + dev = qml.device(dev_name, wires=1) + @qml.qnode(dev) def circuit(x): qml.RX(x, wires=0) diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py index 6780cfe21a1..b20b5cc94bf 100644 --- a/tests/ops/qubit/test_parametric_ops.py +++ b/tests/ops/qubit/test_parametric_ops.py @@ -3464,9 +3464,9 @@ def test_simplify_rotations_grad_jax(self, op): import jax import jax.numpy as jnp - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(simplify, wires, *params, **hyperparams): if simplify: qml.simplify(op(*params, wires=wires, **hyperparams)) diff --git a/tests/ops/qubit/test_qchem_ops.py b/tests/ops/qubit/test_qchem_ops.py index ce3be78b948..f652da4c5a0 100644 --- a/tests/ops/qubit/test_qchem_ops.py +++ b/tests/ops/qubit/test_qchem_ops.py @@ -273,10 +273,10 @@ def test_autograd(self, excitation): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") state = np.array([0, -1 / np.sqrt(2), 1 / np.sqrt(2), 0]) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -298,9 +298,9 @@ def test_autograd_grad(self, diff_method, excitation, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="autograd") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -324,9 +324,9 @@ def test_tf(self, excitation, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -355,9 +355,9 @@ def test_jax(self, excitation, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="jax") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -505,12 +505,12 @@ def test_autograd(self, excitation): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -528,12 +528,12 @@ def test_tf(self, excitation): """Tests that operations are computed correctly using the tensorflow interface""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -551,12 +551,12 @@ def test_jax(self, excitation): """Tests that operations are computed correctly using the jax interface""" - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -579,9 +579,9 @@ def test_autograd_grad(self, excitation, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -607,9 +607,9 @@ def test_tf_grad(self, excitation, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -639,9 +639,9 @@ def test_jax_grad(self, excitation, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -774,7 +774,7 @@ def test_autograd(self): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -796,7 +796,7 @@ def test_autograd(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -811,7 +811,7 @@ def test_tf(self): """Tests that operations are computed correctly using the tensorflow interface""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -833,7 +833,7 @@ def test_tf(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -848,7 +848,7 @@ def test_jax(self): """Tests that operations are computed correctly using the jax interface""" - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -870,7 +870,7 @@ def test_jax(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -885,7 +885,7 @@ def test_torch(self): """Tests that operations are computed correctly using the torch interface""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -907,7 +907,7 @@ def test_torch(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="torch") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -930,10 +930,14 @@ def test_autograd_grad(self, phi, diff_method): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode( + self.grad_circuit_0, dev, interface="autograd", diff_method=diff_method + ) + circuit_1 = qml.QNode( + self.grad_circuit_1, dev, interface="autograd", diff_method=diff_method + ) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) assert np.allclose(qml.grad(total)(phi), self.expected_grad_fn(phi)) @@ -950,10 +954,10 @@ def test_tf_grad(self, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="tf", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="tf", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_t = tf.Variable(phi, dtype=tf.float64) @@ -976,10 +980,10 @@ def test_jax_grad(self, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="jax", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="jax", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_j = jax.numpy.array(phi) @@ -998,10 +1002,10 @@ def test_torch_grad(self, phi, diff_method): import torch - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="torch", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="torch", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_t = torch.tensor(phi, dtype=torch.complex128, requires_grad=True) @@ -1105,7 +1109,7 @@ def test_autograd(self): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") state = np.array( [ 0, @@ -1115,7 +1119,7 @@ def test_autograd(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1137,9 +1141,9 @@ def test_autograd_grad(self, diff_method, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="autograd", diff_method=diff_method) def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1163,9 +1167,9 @@ def test_tf(self, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1194,9 +1198,9 @@ def test_jax(self, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="jax", diff_method=diff_method) def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) diff --git a/tests/optimize/test_spsa.py b/tests/optimize/test_spsa.py index 01726f843f1..f0422602ffd 100644 --- a/tests/optimize/test_spsa.py +++ b/tests/optimize/test_spsa.py @@ -443,7 +443,7 @@ def cost(params): @pytest.mark.usefixtures("use_legacy_opmath") @pytest.mark.slow - def test_lighting_device_legacy_opmath(self): + def test_lightning_device_legacy_opmath(self): """Test SPSAOptimizer implementation with lightning.qubit device.""" coeffs = [0.2, -0.543, 0.4514] obs = [ @@ -479,7 +479,7 @@ def cost_fun(params, num_qubits=1): assert energy < init_energy @pytest.mark.slow - def test_lighting_device(self): + def test_lightning_device(self): """Test SPSAOptimizer implementation with lightning.qubit device.""" coeffs = [0.2, -0.543, 0.4514] obs = [ @@ -494,6 +494,9 @@ def test_lighting_device(self): @qml.qnode(dev) def cost_fun(params, num_qubits=1): qml.BasisState([1, 1, 0, 0], wires=range(num_qubits)) + + assert num_qubits == 4 + for i in range(num_qubits): qml.Rot(*params[i], wires=0) qml.CNOT(wires=[2, 3]) diff --git a/tests/pulse/test_parametrized_evolution.py b/tests/pulse/test_parametrized_evolution.py index 7888d71673a..2fced214278 100644 --- a/tests/pulse/test_parametrized_evolution.py +++ b/tests/pulse/test_parametrized_evolution.py @@ -155,7 +155,6 @@ def test_init(self, params, coeffs): assert ev.num_wires == AnyWires assert ev.name == "ParametrizedEvolution" assert ev.id is None - assert ev.queue_idx is None exp_params = [] if params is None else params assert qml.math.allequal(ev.data, exp_params) diff --git a/tests/pulse/test_transmon.py b/tests/pulse/test_transmon.py index eb280f04bab..6a0631d3abc 100644 --- a/tests/pulse/test_transmon.py +++ b/tests/pulse/test_transmon.py @@ -510,7 +510,7 @@ def fb(p, t): Hd = transmon_drive(amplitude=fa, phase=fb, freq=0.5, wires=[0]) H = Hi + Hd - dev = qml.device("default.qubit.jax", wires=wires) + dev = qml.device("default.qubit") ts = jnp.array([0.0, 3.0]) H_obj = sum(qml.PauliZ(i) for i in range(2)) diff --git a/tests/qchem/test_hamiltonians.py b/tests/qchem/test_hamiltonians.py index 1f71a8151ca..26c640b3f8e 100644 --- a/tests/qchem/test_hamiltonians.py +++ b/tests/qchem/test_hamiltonians.py @@ -267,6 +267,36 @@ def test_diff_hamiltonian_active_space(): assert isinstance(h, qml.ops.Sum if active_new_opmath() else qml.Hamiltonian) +@pytest.mark.parametrize( + ("symbols", "geometry", "core", "active", "charge"), + [ + ( + ["H", "H"], + np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 2.0]]), + None, + None, + 0, + ), + ( + ["H", "H", "H"], + np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 1.0], [0.0, 2.0, 0.0]]), + [0], + [1, 2], + 1, + ), + ], +) +def test_diff_hamiltonian_wire_order(symbols, geometry, core, active, charge): + r"""Test that diff_hamiltonian has an ascending wire order.""" + + mol = qchem.Molecule(symbols, geometry, charge) + args = [geometry] + + h = qchem.diff_hamiltonian(mol, core=core, active=active)(*args) + + assert h.wires.tolist() == sorted(h.wires.tolist()) + + def test_gradient_expvalH(): r"""Test that the gradient of expval(H) computed with ``qml.grad`` is equal to the value obtained with the finite difference method.""" diff --git a/tests/qnn/test_keras.py b/tests/qnn/test_keras.py index af8d457cd9e..f4f9769edc2 100644 --- a/tests/qnn/test_keras.py +++ b/tests/qnn/test_keras.py @@ -503,7 +503,7 @@ def test_gradients(self, get_circuit, output_dim, n_qubits): # pylint: disable= def test_backprop_gradients(self, mocker): # pylint: disable=no-self-use """Test if KerasLayer is compatible with the backprop diff method.""" - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") @qml.qnode(dev, interface="tf", diff_method="backprop") def f(inputs, weights): @@ -813,9 +813,9 @@ def test_no_attribute(): @pytest.mark.tf def test_batch_input_single_measure(tol): """Test input batching in keras""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="tf", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -830,7 +830,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 2) - assert dev.num_executions == 1 for x_, r in zip(x, res): assert qml.math.allclose(r, circuit(x_, layer.qnode_weights["weights"]), atol=tol) @@ -842,9 +841,9 @@ def circuit(x, weights): @pytest.mark.tf def test_batch_input_multi_measure(tol): """Test input batching in keras for multiple measurements""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="tf", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -859,7 +858,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 5) - assert dev.num_executions == 1 for x_, r in zip(x, res): exp = tf.experimental.numpy.hstack(circuit(x_, layer.qnode_weights["weights"])) diff --git a/tests/qnn/test_qnn_torch.py b/tests/qnn/test_qnn_torch.py index a25ee7d6949..64aeb9b1a9c 100644 --- a/tests/qnn/test_qnn_torch.py +++ b/tests/qnn/test_qnn_torch.py @@ -805,9 +805,9 @@ def circ(inputs, w0): # pylint: disable=unused-argument @pytest.mark.torch def test_batch_input_single_measure(tol): """Test input batching in torch""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="torch", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -820,7 +820,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 2) - assert dev.num_executions == 1 for x_, r in zip(x, res): assert qml.math.allclose(r, circuit(x_, layer.qnode_weights["weights"]), atol=tol) @@ -832,9 +831,9 @@ def circuit(x, weights): @pytest.mark.torch def test_batch_input_multi_measure(tol): """Test input batching in torch for multiple measurements""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="torch", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -847,7 +846,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 5) - assert dev.num_executions == 1 for x_, r in zip(x, res): exp = torch.hstack(circuit(x_, layer.qnode_weights["weights"])) diff --git a/tests/resource/test_specs.py b/tests/resource/test_specs.py index 2dd69be8dd3..ac2744d50e1 100644 --- a/tests/resource/test_specs.py +++ b/tests/resource/test_specs.py @@ -14,12 +14,20 @@ """Unit tests for the specs transform""" from collections import defaultdict from contextlib import nullcontext -from typing import Callable, Sequence import pytest import pennylane as qml from pennylane import numpy as pnp +from pennylane.tape import QuantumTapeBatch +from pennylane.typing import PostprocessingFn + +with pytest.warns(qml.PennyLaneDeprecationWarning): + devices_list = [ + (qml.device("default.qubit"), 1), + (qml.device("default.qubit", wires=2), 2), + (qml.device("default.qubit.legacy", wires=2), 2), + ] class TestSpecsTransform: @@ -48,7 +56,7 @@ def test_max_expansion_throws_warning(self): def circ(): return qml.expval(qml.PauliZ(0)) - with pytest.warns(UserWarning, match="'max_expansion' has no effect"): + with pytest.warns(qml.PennyLaneDeprecationWarning, match="'max_expansion' has no effect"): qml.specs(circ, max_expansion=10)() def test_only_one_of_level_or_expansion_strategy_passed(self): @@ -203,7 +211,7 @@ def test_splitting_transforms(self): obs = [qml.X(0) @ qml.Z(1), qml.Z(0) @ qml.Y(2), qml.Y(0) @ qml.X(2)] H = qml.Hamiltonian(coeffs, obs) - @qml.transforms.hamiltonian_expand + @qml.transforms.split_non_commuting @qml.transforms.merge_rotations @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=pnp.pi / 4) def circuit(x): @@ -251,13 +259,16 @@ def circuit(params): @pytest.mark.xfail(reason="DefaultQubit2 does not support custom expansion depths") def test_max_expansion(self): - """Test that a user can calculation specifications for a different max + """Test that a user can calculate specifications for a different max expansion parameter.""" circuit, params = self.make_qnode_and_params("device") assert circuit.max_expansion == 10 - info = qml.specs(circuit, max_expansion=0)(params) + + with pytest.warns(UserWarning, match="'max_expansion' has no effect"): + info = qml.specs(circuit, max_expansion=0)(params) + assert circuit.max_expansion == 10 assert len(info) == 11 @@ -276,10 +287,18 @@ def test_max_expansion(self): def test_expansion_strategy(self): """Test that a user can calculate specs for different expansion strategies.""" - circuit, params = self.make_qnode_and_params("gradient") + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' attribute is deprecated" + ): + circuit, params = self.make_qnode_and_params("gradient") assert circuit.expansion_strategy == "gradient" - info = qml.specs(circuit, expansion_strategy="device")(params) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="'expansion_strategy' argument is deprecated" + ): + info = qml.specs(circuit, expansion_strategy="device")(params) + assert circuit.expansion_strategy == "gradient" assert len(info) == 13 @@ -304,7 +323,7 @@ def test_custom_gradient_transform(self): @qml.transforms.core.transform def my_transform( tape: qml.tape.QuantumTape, - ) -> tuple[Sequence[qml.tape.QuantumTape], Callable]: + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return tape, None @qml.qnode(dev, diff_method=my_transform) @@ -317,11 +336,7 @@ def circuit(): @pytest.mark.parametrize( "device,num_wires", - [ - (qml.device("default.qubit"), 1), - (qml.device("default.qubit", wires=2), 2), - (qml.device("default.qubit.legacy", wires=2), 2), - ], + devices_list, ) def test_num_wires_source_of_truth(self, device, num_wires): """Tests that num_wires behaves differently on old and new devices.""" diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py index 5cf0f47c899..859feb114d3 100644 --- a/tests/templates/test_embeddings/test_amplitude.py +++ b/tests/templates/test_embeddings/test_amplitude.py @@ -59,7 +59,7 @@ def test_expansion(self): """Checks the queue for the default settings.""" op = qml.AmplitudeEmbedding(features=FEATURES[0], wires=range(2)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == 1 assert tape.operations[0].name == "StatePrep" @@ -70,7 +70,7 @@ def test_expansion_broadcasted(self): op = qml.AmplitudeEmbedding(features=BROADCASTED_FEATURES[0], wires=range(2)) assert op.batch_size == 3 - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == 1 assert tape.operations[0].name == "StatePrep" diff --git a/tests/templates/test_embeddings/test_angle.py b/tests/templates/test_embeddings/test_angle.py index 2d7bf5b0f68..93cfdaccc0b 100644 --- a/tests/templates/test_embeddings/test_angle.py +++ b/tests/templates/test_embeddings/test_angle.py @@ -59,7 +59,7 @@ def test_expansion(self, features): """Checks the queue for the default settings.""" op = qml.AngleEmbedding(features=features, wires=range(4)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == len(features) for gate in tape.operations: @@ -73,7 +73,7 @@ def test_expansion_broadcasted(self): op = qml.AngleEmbedding(features=features, wires=range(4)) assert op.batch_size == 5 - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == 3 for gate in tape.operations: @@ -86,7 +86,7 @@ def test_rotations(self, rotation): """Checks the queue for the specified rotation settings.""" op = qml.AngleEmbedding(features=[1, 1, 1], wires=range(4), rotation=rotation) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for gate in tape.operations: assert gate.name == "R" + rotation diff --git a/tests/templates/test_embeddings/test_basis.py b/tests/templates/test_embeddings/test_basis.py index 81ea1b52616..31bc736cff6 100644 --- a/tests/templates/test_embeddings/test_basis.py +++ b/tests/templates/test_embeddings/test_basis.py @@ -54,7 +54,7 @@ def test_expansion(self, features): """Checks the queue.""" op = qml.BasisEmbedding(features=features, wires=range(3)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == features.count(1) for gate in tape.operations: diff --git a/tests/templates/test_embeddings/test_displacement_emb.py b/tests/templates/test_embeddings/test_displacement_emb.py index 2700df62e72..9c9769ab469 100644 --- a/tests/templates/test_embeddings/test_displacement_emb.py +++ b/tests/templates/test_embeddings/test_displacement_emb.py @@ -54,7 +54,7 @@ def test_expansion(self, features): """Checks the queue for the default settings.""" op = qml.DisplacementEmbedding(features=features, wires=range(3)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == len(features) for idx, gate in enumerate(tape.operations): diff --git a/tests/templates/test_embeddings/test_iqp_emb.py b/tests/templates/test_embeddings/test_iqp_emb.py index c19731c77fc..274e9e3e7a5 100644 --- a/tests/templates/test_embeddings/test_iqp_emb.py +++ b/tests/templates/test_embeddings/test_iqp_emb.py @@ -49,7 +49,7 @@ def test_expansion(self, n_wires, expected_names, expected_wires): features = list(range(n_wires)) op = qml.IQPEmbedding(features, wires=range(n_wires)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) j = 0 for i, gate in enumerate(tape.operations): @@ -67,7 +67,7 @@ def test_expansion_broadcasted(self, n_wires, expected_names, expected_wires): op = qml.IQPEmbedding(features, wires=range(n_wires)) assert op.batch_size == 3 - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) j = 0 for i, gate in enumerate(tape.operations): @@ -87,7 +87,7 @@ def test_repeat(self): expected_wires = self.QUEUES[2][2] + self.QUEUES[2][2] op = qml.IQPEmbedding(features, wires=range(3), n_repeats=2) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -112,7 +112,7 @@ def test_custom_pattern(self): expected_wires = [[0], [0], [1], [1], [2], [2], *pattern] op = qml.IQPEmbedding(features, wires=range(3), pattern=pattern) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] diff --git a/tests/templates/test_embeddings/test_qaoa_emb.py b/tests/templates/test_embeddings/test_qaoa_emb.py index 92909c35f2f..4a56f14b15b 100644 --- a/tests/templates/test_embeddings/test_qaoa_emb.py +++ b/tests/templates/test_embeddings/test_qaoa_emb.py @@ -76,7 +76,7 @@ def test_expansion(self, n_wires, weight_shape, expected_names): weights = np.zeros(shape=weight_shape) op = qml.QAOAEmbedding(features, weights, wires=range(n_wires)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -94,7 +94,7 @@ def test_expansion_broadcasted(self, n_wires, weight_shape, expected_names): # Only broadcast features op = qml.QAOAEmbedding(broadcasted_features, weights, wires=range(n_wires)) assert op.batch_size == n_broadcast - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -106,7 +106,7 @@ def test_expansion_broadcasted(self, n_wires, weight_shape, expected_names): # Only broadcast weights op = qml.QAOAEmbedding(features, broadcasted_weights, wires=range(n_wires)) assert op.batch_size == n_broadcast - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -118,7 +118,7 @@ def test_expansion_broadcasted(self, n_wires, weight_shape, expected_names): # Broadcast weights and features op = qml.QAOAEmbedding(broadcasted_features, broadcasted_weights, wires=range(n_wires)) assert op.batch_size == n_broadcast - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -134,7 +134,7 @@ def test_local_field(self, local_field): weights = np.zeros(shape=(1, 3)) op = qml.QAOAEmbedding(features, weights, wires=range(2), local_field=local_field) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) gate_names = [gate.name for gate in tape.operations] assert gate_names[3] == get_name[local_field] diff --git a/tests/templates/test_embeddings/test_squeezing_emb.py b/tests/templates/test_embeddings/test_squeezing_emb.py index 4486ee7cde7..198905e6614 100644 --- a/tests/templates/test_embeddings/test_squeezing_emb.py +++ b/tests/templates/test_embeddings/test_squeezing_emb.py @@ -54,7 +54,7 @@ def test_expansion(self, features): """Checks the queue for the default settings.""" op = qml.SqueezingEmbedding(features=features, wires=range(3)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) assert len(tape.operations) == len(features) for idx, gate in enumerate(tape.operations): diff --git a/tests/templates/test_layers/test_basic_entangler.py b/tests/templates/test_layers/test_basic_entangler.py index 4bce9032250..3f672917b2d 100644 --- a/tests/templates/test_layers/test_basic_entangler.py +++ b/tests/templates/test_layers/test_basic_entangler.py @@ -54,7 +54,7 @@ def test_expansion(self, n_wires, weight_shape, expected_names, expected_wires): weights = np.random.random(size=weight_shape) op = qml.BasicEntanglerLayers(weights, wires=range(n_wires)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) for i, gate in enumerate(tape.operations): assert gate.name == expected_names[i] @@ -67,7 +67,7 @@ def test_rotation(self, rotation): weights = np.zeros(shape=(1, 2)) op = qml.BasicEntanglerLayers(weights, wires=range(2), rotation=rotation) - queue = op.expand().operations + queue = op.decomposition() assert rotation in [type(gate) for gate in queue] diff --git a/tests/templates/test_layers/test_cv_neural_net.py b/tests/templates/test_layers/test_cv_neural_net.py index ddb4c35a2c2..d9946cf6740 100644 --- a/tests/templates/test_layers/test_cv_neural_net.py +++ b/tests/templates/test_layers/test_cv_neural_net.py @@ -76,7 +76,7 @@ def test_expansion(self, n_wires, expected_names, expected_wires): weights = [np.random.random(shape) for shape in shapes] op = qml.CVNeuralNetLayers(*weights, wires=range(n_wires)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) i = 0 for gate in tape.operations: @@ -85,7 +85,7 @@ def test_expansion(self, n_wires, expected_names, expected_wires): assert gate.wires.labels == tuple(expected_wires[i]) i = i + 1 else: - for gate_inter in gate.expand().operations: + for gate_inter in gate.decomposition(): assert gate_inter.name == expected_names[i] assert gate_inter.wires.labels == tuple(expected_wires[i]) i = i + 1 diff --git a/tests/templates/test_layers/test_gate_fabric.py b/tests/templates/test_layers/test_gate_fabric.py index ae14dc173a6..a69a8fe42a0 100644 --- a/tests/templates/test_layers/test_gate_fabric.py +++ b/tests/templates/test_layers/test_gate_fabric.py @@ -72,7 +72,7 @@ def test_operations(self, layers, qubits, init_state, include_pi): op = qml.GateFabric( weights, wires=range(qubits), init_state=init_state, include_pi=include_pi ) - queue = op.expand().operations + queue = op.decomposition() # number of gates assert len(queue) == n_gates diff --git a/tests/templates/test_layers/test_particle_conserving_u1.py b/tests/templates/test_layers/test_particle_conserving_u1.py index 0f9c23cdc13..7e46abe879f 100644 --- a/tests/templates/test_layers/test_particle_conserving_u1.py +++ b/tests/templates/test_layers/test_particle_conserving_u1.py @@ -90,7 +90,7 @@ def test_operations(self): nm_wires += [wires[l : l + 2] for l in range(1, qubits - 1, 2)] op = qml.ParticleConservingU1(weights, wires, init_state=np.array([1, 1, 0, 0])) - queue = op.expand().operations + queue = op.decomposition() assert gate_count == len(queue) diff --git a/tests/templates/test_layers/test_particle_conserving_u2.py b/tests/templates/test_layers/test_particle_conserving_u2.py index 54e92914a3f..c6c88bbfe88 100644 --- a/tests/templates/test_layers/test_particle_conserving_u2.py +++ b/tests/templates/test_layers/test_particle_conserving_u2.py @@ -57,7 +57,7 @@ def test_operations(self, layers, qubits, init_state): ) * layers op = qml.ParticleConservingU2(weights, wires=range(qubits), init_state=init_state) - queue = op.expand().operations + queue = op.decomposition() # number of gates assert len(queue) == n_gates diff --git a/tests/templates/test_layers/test_random.py b/tests/templates/test_layers/test_random.py index 42a7accc398..9d976307426 100644 --- a/tests/templates/test_layers/test_random.py +++ b/tests/templates/test_layers/test_random.py @@ -83,11 +83,11 @@ def test_seed(self): op2 = qml.RandomLayers(weights, wires=range(2), seed=42) op3 = qml.RandomLayers(weights, wires=range(2), seed=42) - queue1 = op1.expand().operations + queue1 = op1.decomposition() decomp1 = op1.compute_decomposition(*op1.parameters, wires=op1.wires, **op1.hyperparameters) - queue2 = op2.expand().operations + queue2 = op2.decomposition() decomp2 = op2.compute_decomposition(*op2.parameters, wires=op2.wires, **op2.hyperparameters) - queue3 = op3.expand().operations + queue3 = op3.decomposition() decomp3 = op3.compute_decomposition(*op3.parameters, wires=op3.wires, **op3.hyperparameters) assert not all(g1.name == g2.name for g1, g2 in zip(queue1, queue2)) @@ -106,7 +106,7 @@ def test_number_gates(self, n_layers, n_rots): weights = np.random.randn(n_layers, n_rots) op = qml.RandomLayers(weights, wires=range(2)) - ops = op.expand().operations + ops = op.decomposition() gate_names = [g.name for g in ops] assert len(gate_names) - gate_names.count("CNOT") == n_layers * n_rots @@ -131,7 +131,7 @@ def test_random_wires(self): weights = np.random.random(size=(2, n_rots)) op = qml.RandomLayers(weights, wires=range(3)) - queue = op.expand().operations + queue = op.decomposition() gate_wires = [gate.wires.labels for gate in queue] wires_flat = [item for w in gate_wires for item in w] diff --git a/tests/templates/test_layers/test_simplified_twodesign.py b/tests/templates/test_layers/test_simplified_twodesign.py index abe9ca025bd..1580fb42a25 100644 --- a/tests/templates/test_layers/test_simplified_twodesign.py +++ b/tests/templates/test_layers/test_simplified_twodesign.py @@ -64,7 +64,7 @@ def test_expansion(self, n_wires, weight_shape, expected_names, expected_wires): initial_layer = np.random.randn(n_wires) op = qml.SimplifiedTwoDesign(initial_layer, weights, wires=range(n_wires)) - queue = op.expand().operations + queue = op.decomposition() for i, gate in enumerate(queue): assert gate.name == expected_names[i] @@ -81,7 +81,7 @@ def test_circuit_parameters(self, n_wires, n_layers, shape_weights): weights = np.random.randn(*shape_weights) op = qml.SimplifiedTwoDesign(initial_layer, weights, wires=range(n_wires)) - queue = op.expand().operations + queue = op.decomposition() # test the device parameters for _ in range(n_layers): diff --git a/tests/templates/test_layers/test_strongly_entangling.py b/tests/templates/test_layers/test_strongly_entangling.py index ed9839b840b..393786b3d2d 100644 --- a/tests/templates/test_layers/test_strongly_entangling.py +++ b/tests/templates/test_layers/test_strongly_entangling.py @@ -67,7 +67,7 @@ def test_expansion(self, n_wires, weight_shape, expected_names, expected_wires, weights = np.random.random(size=weight_shape) op = qml.StronglyEntanglingLayers(weights, wires=range(n_wires)) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) if batch_dim is None: param_sets = iter(weights.reshape((-1, 3))) @@ -93,7 +93,7 @@ def test_uses_correct_imprimitive(self, n_layers, n_wires, batch_dim): weights = np.random.randn(*shape) op = qml.StronglyEntanglingLayers(weights=weights, wires=range(n_wires), imprimitive=qml.CZ) - ops = op.expand().operations + ops = op.decomposition() gate_names = [gate.name for gate in ops] assert gate_names.count("CZ") == n_wires * n_layers @@ -134,7 +134,7 @@ def test_custom_range_sequence(self, n_layers, n_wires, ranges, batch_dim): weights = np.random.randn(*shape) op = qml.StronglyEntanglingLayers(weights=weights, wires=range(n_wires), ranges=ranges) - ops = op.expand().operations + ops = op.decomposition() gate_wires = [gate.wires.labels for gate in ops] range_idx = 0 diff --git a/tests/templates/test_state_preparations/test_arbitrary_state_prep.py b/tests/templates/test_state_preparations/test_arbitrary_state_prep.py index 8191dbc5f0c..e12f1dc5f9d 100644 --- a/tests/templates/test_state_preparations/test_arbitrary_state_prep.py +++ b/tests/templates/test_state_preparations/test_arbitrary_state_prep.py @@ -79,7 +79,7 @@ def test_correct_gates_single_wire(self): weights = np.array([0, 1], dtype=float) op = qml.ArbitraryStatePreparation(weights, wires=[0]) - queue = op.expand().operations + queue = op.decomposition() assert queue[0].name == "PauliRot" @@ -97,7 +97,7 @@ def test_correct_gates_two_wires(self): weights = np.array([0, 1, 2, 3, 4, 5], dtype=float) op = qml.ArbitraryStatePreparation(weights, wires=[0, 1]) - queue = op.expand().operations + queue = op.decomposition() assert queue[0].name == "PauliRot" diff --git a/tests/templates/test_state_preparations/test_basis_state_prep.py b/tests/templates/test_state_preparations/test_basis_state_prep.py index f7dcd2a2e56..90b88463a05 100644 --- a/tests/templates/test_state_preparations/test_basis_state_prep.py +++ b/tests/templates/test_state_preparations/test_basis_state_prep.py @@ -55,7 +55,7 @@ def test_correct_pl_gates(self, basis_state, wires, target_wires): """Tests queue for simple cases.""" op = qml.BasisStatePreparation(basis_state, wires) - queue = op.expand().operations + queue = op.decomposition() for id, gate in enumerate(queue): assert gate.name == "PauliX" diff --git a/tests/templates/test_state_preparations/test_cosine_window.py b/tests/templates/test_state_preparations/test_cosine_window.py index 2941d9f6b1e..a9bfa87fc83 100644 --- a/tests/templates/test_state_preparations/test_cosine_window.py +++ b/tests/templates/test_state_preparations/test_cosine_window.py @@ -38,7 +38,7 @@ def test_correct_gates_single_wire(self): """Test that the correct gates are applied.""" op = qml.CosineWindow(wires=[0]) - queue = op.expand().operations + queue = op.decomposition() assert queue[0].name == "Hadamard" assert queue[1].name == "RZ" diff --git a/tests/templates/test_state_preparations/test_qutrit_basis_state_prep.py b/tests/templates/test_state_preparations/test_qutrit_basis_state_prep.py index c52d5d334a3..98755079512 100644 --- a/tests/templates/test_state_preparations/test_qutrit_basis_state_prep.py +++ b/tests/templates/test_state_preparations/test_qutrit_basis_state_prep.py @@ -58,7 +58,7 @@ def test_correct_pl_gates(self, basis_state, wires, target_wires): """Tests queue for simple cases.""" op = qml.QutritBasisStatePreparation(basis_state, wires) - queue = op.expand().operations + queue = op.decomposition() for id, gate in enumerate(queue): assert gate.name == "TShift" diff --git a/tests/templates/test_subroutines/test_all_singles_doubles.py b/tests/templates/test_subroutines/test_all_singles_doubles.py index b303beab5a3..7d20863135d 100644 --- a/tests/templates/test_subroutines/test_all_singles_doubles.py +++ b/tests/templates/test_subroutines/test_all_singles_doubles.py @@ -99,7 +99,7 @@ def test_allsinglesdoubles_operations(self, singles, doubles, weights, ref_gates hf_state = np.array([1, 1, 0, 0, 0, 0]) op = qml.AllSinglesDoubles(weights, wires, hf_state, singles=singles, doubles=doubles) - queue = op.expand().operations + queue = op.decomposition() assert len(queue) == len(singles) + len(doubles) + 1 diff --git a/tests/templates/test_subroutines/test_approx_time_evolution.py b/tests/templates/test_subroutines/test_approx_time_evolution.py index 80f6b962b98..42623312625 100644 --- a/tests/templates/test_subroutines/test_approx_time_evolution.py +++ b/tests/templates/test_subroutines/test_approx_time_evolution.py @@ -116,7 +116,7 @@ def test_evolution_operations(self, time, hamiltonian, steps, expected_queue): """Tests that the sequence of gates implemented in the ApproxTimeEvolution template is correct""" op = qml.ApproxTimeEvolution(hamiltonian, time, steps) - queue = op.expand().operations + queue = op.decomposition() for expected_gate, gate in zip(expected_queue, queue): qml.assert_equal(expected_gate, gate) diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py index 9578c75f655..c5c9c643450 100644 --- a/tests/templates/test_subroutines/test_basis_rotation.py +++ b/tests/templates/test_subroutines/test_basis_rotation.py @@ -86,7 +86,7 @@ def test_basis_rotation_operations(self, num_wires, unitary_matrix, givens, diag gate_wires.append(list(indices)) op = qml.BasisRotation(wires=range(num_wires), unitary_matrix=unitary_matrix) - queue = op.expand().operations + queue = op.decomposition() assert len(queue) == len(gate_ops) # number of gates diff --git a/tests/templates/test_subroutines/test_commuting_evolution.py b/tests/templates/test_subroutines/test_commuting_evolution.py index a578e2d19f2..96950dcc265 100644 --- a/tests/templates/test_subroutines/test_commuting_evolution.py +++ b/tests/templates/test_subroutines/test_commuting_evolution.py @@ -89,7 +89,7 @@ def test_decomposition_expand(): assert qml.math.allclose(decomp.hyperparameters["hamiltonian"].data, hamiltonian.data) assert decomp.hyperparameters["n"] == 1 - tape = op.expand() + tape = op.decomposition() assert len(tape) == 1 assert isinstance(tape[0], qml.ApproxTimeEvolution) diff --git a/tests/templates/test_subroutines/test_double_excitation.py b/tests/templates/test_subroutines/test_double_excitation.py index c55a6db4921..d299b7ac942 100644 --- a/tests/templates/test_subroutines/test_double_excitation.py +++ b/tests/templates/test_subroutines/test_double_excitation.py @@ -203,7 +203,7 @@ def test_double_ex_unitary_operations(self, wires1, wires2, ref_gates): cnots = 16 * (len(wires1) - 1 + len(wires2) - 1 + 1) weight = np.pi / 3 op = qml.FermionicDoubleExcitation(weight, wires1=wires1, wires2=wires2) - queue = op.expand().operations + queue = op.decomposition() assert len(queue) == sqg + cnots diff --git a/tests/templates/test_subroutines/test_grover.py b/tests/templates/test_subroutines/test_grover.py index 3bb39e11628..3ff3755a87a 100644 --- a/tests/templates/test_subroutines/test_grover.py +++ b/tests/templates/test_subroutines/test_grover.py @@ -47,7 +47,7 @@ def test_work_wires(): assert op.hyperparameters["work_wires"] == work_wire - ops = op.expand().operations + ops = op.decomposition() assert ops[2].hyperparameters["work_wires"] == work_wire @@ -191,7 +191,7 @@ def test_expand(wires): """Asserts decomposition uses expected operations and wires""" op = qml.GroverOperator(wires=wires) - decomp = op.expand().operations + decomp = op.decomposition() expected_wires = decomposition_wires(wires) diff --git a/tests/templates/test_subroutines/test_interferometer.py b/tests/templates/test_subroutines/test_interferometer.py index 9f31ab9895e..6221d6f69eb 100644 --- a/tests/templates/test_subroutines/test_interferometer.py +++ b/tests/templates/test_subroutines/test_interferometer.py @@ -70,7 +70,7 @@ def test_clements_beamsplitter_convention(self): theta, phi, varphi, mesh="triangular", beamsplitter="clements", wires=wires ) - for rec in [op_rect.expand(), op_tria.expand()]: + for rec in [op_rect.decomposition(), op_tria.decomposition()]: assert len(rec) == 4 assert isinstance(rec[0], qml.Rotation) @@ -90,7 +90,7 @@ def test_one_mode(self, tol): varphi = [0.42342] op = qml.Interferometer(theta=[], phi=[], varphi=varphi, wires=0) - rec = op.expand() + rec = op.decomposition() assert len(rec) == 1 assert isinstance(rec[0], qml.Rotation) @@ -107,7 +107,7 @@ def test_two_mode_rect(self): varphi = [0.42342, 0.1121] op = qml.Interferometer(theta, phi, varphi, wires=wires) - rec = op.expand() + rec = op.decomposition() isinstance(rec[0], qml.Beamsplitter) assert rec[0].parameters == theta + phi @@ -129,7 +129,7 @@ def test_two_mode_triangular(self): varphi = [0.42342, 0.1121] op = qml.Interferometer(theta, phi, varphi, mesh="triangular", wires=wires) - rec = op.expand() + rec = op.decomposition() assert len(rec) == 3 @@ -155,7 +155,7 @@ def test_three_mode(self): op_tria = qml.Interferometer(theta, phi, varphi, wires=wires, mesh="triangular") # Test rectangular mesh - rec = op_rect.expand() + rec = op_rect.decomposition() assert len(rec) == 6 expected_bs_wires = [[0, 1], [1, 2], [0, 1]] @@ -171,7 +171,7 @@ def test_three_mode(self): assert op.wires == Wires([idx]) # Test triangular mesh - rec = op_tria.expand() + rec = op_tria.decomposition() assert len(rec) == 6 expected_bs_wires = [[1, 2], [0, 1], [1, 2]] @@ -196,7 +196,7 @@ def test_four_mode_rect(self): varphi = [0.42342, 0.234, 0.4523, 0.1121] op = qml.Interferometer(theta, phi, varphi, wires=wires) - rec = op.expand() + rec = op.decomposition() assert len(rec) == 10 @@ -222,7 +222,7 @@ def test_four_mode_triangular(self): varphi = [0.42342, 0.234, 0.4523, 0.1121] op = qml.Interferometer(theta, phi, varphi, wires=wires, mesh="triangular") - rec = op.expand() + rec = op.decomposition() assert len(rec) == 10 diff --git a/tests/templates/test_subroutines/test_kupccgsd.py b/tests/templates/test_subroutines/test_kupccgsd.py index 17e9eddbdfd..ca1ce3f6dbc 100644 --- a/tests/templates/test_subroutines/test_kupccgsd.py +++ b/tests/templates/test_subroutines/test_kupccgsd.py @@ -91,7 +91,7 @@ def test_kupccgsd_operations(self, k, delta_sz, init_state, wires): exp_unitary += [qml.FermionicSingleExcitation] * len(gen_single_terms_wires) op = qml.kUpCCGSD(weights, wires=wires, k=k, delta_sz=delta_sz, init_state=init_state) - queue = op.expand().operations + queue = op.decomposition() # number of gates assert len(queue) == n_gates diff --git a/tests/templates/test_subroutines/test_prepselprep.py b/tests/templates/test_subroutines/test_prepselprep.py index 587bbdccb26..82629973865 100644 --- a/tests/templates/test_subroutines/test_prepselprep.py +++ b/tests/templates/test_subroutines/test_prepselprep.py @@ -27,6 +27,10 @@ ("lcu", "control"), [ (qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]), + (qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]), + (qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]), + (0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2), [0]), + (qml.Z(2) + qml.X(1) @ qml.X(2), [0]), (qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]), [0, 1]), ( qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]), @@ -257,25 +261,35 @@ def test_copy(self): assert qml.equal(op, op_copy) - def test_flatten_unflatten(self): + @pytest.mark.parametrize( + ("lcu"), + [ + qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), + qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), + qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), + 0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2), + qml.Z(2) + qml.X(1) @ qml.X(2), + qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]), + qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]), + ], + ) + def test_flatten_unflatten(self, lcu): """Test that the class can be correctly flattened and unflattened""" - lcu = qml.ops.LinearCombination([1 / 2, 1 / 2], [qml.Identity(1), qml.PauliZ(1)]) lcu_coeffs, lcu_ops = lcu.terms() op = qml.PrepSelPrep(lcu, control=0) data, metadata = op._flatten() - data_coeffs = [term.terms()[0][0] for term in data] - data_ops = [term.terms()[1][0] for term in data] + data_coeffs, data_ops = data[0].terms() assert hash(metadata) - assert len(data) == len(lcu) + assert len(data[0]) == len(lcu) assert all(coeff1 == coeff2 for coeff1, coeff2 in zip(lcu_coeffs, data_coeffs)) assert all(op1 == op2 for op1, op2 in zip(lcu_ops, data_ops)) - assert metadata == op.control + assert metadata[0] == op.control new_op = type(op)._unflatten(*op._flatten()) assert op.lcu == new_op.lcu diff --git a/tests/templates/test_subroutines/test_qmc.py b/tests/templates/test_subroutines/test_qmc.py index 0353f84b4e5..9a5b18f3c7c 100644 --- a/tests/templates/test_subroutines/test_qmc.py +++ b/tests/templates/test_subroutines/test_qmc.py @@ -297,7 +297,7 @@ def test_expected_circuit(self): target_wires, estimation_wires = Wires(range(3)), Wires(range(3, 5)) op = QuantumMonteCarlo(p, self.func, target_wires, estimation_wires) - tape = op.expand() + tape = qml.tape.QuantumScript(op.decomposition()) # Do expansion in two steps to avoid also decomposing the first QubitUnitary queue_before_qpe = tape.operations[:2] diff --git a/tests/templates/test_subroutines/test_qpe.py b/tests/templates/test_subroutines/test_qpe.py index 7611348a69f..af3d941a85f 100644 --- a/tests/templates/test_subroutines/test_qpe.py +++ b/tests/templates/test_subroutines/test_qpe.py @@ -112,7 +112,7 @@ def test_expected_qscript(self): m = qml.RX(0.3, wires=0).matrix() op = qml.QuantumPhaseEstimation(m, target_wires=[0], estimation_wires=[1, 2]) - qscript = op.expand() + qscript = qml.tape.QuantumScript(op.decomposition()) unitary = qml.QubitUnitary(m, wires=[0]) with qml.queuing.AnnotatedQueue() as q: diff --git a/tests/templates/test_subroutines/test_single_excitation.py b/tests/templates/test_subroutines/test_single_excitation.py index 5d181cd3965..f63b9ef0acd 100644 --- a/tests/templates/test_subroutines/test_single_excitation.py +++ b/tests/templates/test_subroutines/test_single_excitation.py @@ -98,7 +98,7 @@ def test_single_ex_unitary_operations(self, single_wires, ref_gates): cnots = 4 * (len(single_wires) - 1) weight = np.pi / 3 op = qml.FermionicSingleExcitation(weight, wires=single_wires) - queue = op.expand().operations + queue = op.decomposition() assert len(queue) == sqg + cnots diff --git a/tests/templates/test_subroutines/test_uccsd.py b/tests/templates/test_subroutines/test_uccsd.py index 35d86f03dff..f335851bdd6 100644 --- a/tests/templates/test_subroutines/test_uccsd.py +++ b/tests/templates/test_subroutines/test_uccsd.py @@ -190,7 +190,7 @@ def test_uccsd_operations(self, s_wires, d_wires, weights, n_repeats, ref_gates) init_state=ref_state, n_repeats=n_repeats, ) - raw_queue = op.expand().operations + raw_queue = op.decomposition() # hack to avoid updating the test data: # expand the other templates, which now @@ -198,7 +198,7 @@ def test_uccsd_operations(self, s_wires, d_wires, weights, n_repeats, ref_gates) queue = [] for op in raw_queue: if op.name in ["FermionicSingleExcitation", "FermionicDoubleExcitation"]: - queue.extend(op.expand().operations) + queue.extend(op.decomposition()) else: queue.append(op) diff --git a/tests/templates/test_swapnetworks/test_ccl2.py b/tests/templates/test_swapnetworks/test_ccl2.py index b77c324b5b7..e3ebf8e831f 100644 --- a/tests/templates/test_swapnetworks/test_ccl2.py +++ b/tests/templates/test_swapnetworks/test_ccl2.py @@ -80,7 +80,7 @@ def test_ccl2_operations(self, wires, acquaintances, weights, fermionic, shift): op = qml.templates.TwoLocalSwapNetwork( wire_order, acquaintances, weights, fermionic=fermionic, shift=shift ) - queue = op.expand().operations + queue = op.decomposition() # number of gates assert len(queue) == sum( diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 4b66bf2d061..92130d5f3fd 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -353,6 +353,45 @@ def inner(j): assert circuit(5, 6) == 30 # 5 * 6 assert circuit(4, 7) == 28 # 4 * 7 + def test_while_loop_python_fallback(self): + """Test that qml.while_loop fallsback to + Python without qjit""" + + def f(n, m): + @qml.while_loop(lambda i, _: i < n) + def outer(i, sm): + @qml.while_loop(lambda j: j < m) + def inner(j): + return j + 1 + + return i + 1, sm + inner(0) + + return outer(0, 0)[1] + + assert f(5, 6) == 30 # 5 * 6 + assert f(4, 7) == 28 # 4 * 7 + + def test_fallback_while_loop_qnode(self): + """Test that qml.while_loop inside a qnode fallsback to + Python without qjit""" + dev = qml.device("lightning.qubit", wires=1) + + @qml.qnode(dev) + def circuit(n): + @qml.while_loop(lambda v: v[0] < v[1]) + def loop(v): + qml.PauliX(wires=0) + return v[0] + 1, v[1] + + loop((0, n)) + return qml.expval(qml.PauliZ(0)) + + assert jnp.allclose(circuit(1), -1.0) + + res = circuit.tape.operations + expected = [qml.PauliX(0) for i in range(4)] + _ = [qml.assert_equal(i, j) for i, j in zip(res, expected)] + def test_dynamic_wires_for_loops(self): """Test for loops with iteration index-dependant wires.""" dev = qml.device("lightning.qubit", wires=6) @@ -405,6 +444,57 @@ def inner(j): assert jnp.allclose(circuit(4), jnp.eye(2**4)[0]) + def test_for_loop_python_fallback(self): + """Test that qml.for_loop fallsback to Python + interpretation if Catalyst is not available""" + dev = qml.device("lightning.qubit", wires=3) + + @qml.qnode(dev) + def circuit(x, n): + + # for loop with dynamic bounds + @qml.for_loop(0, n, 1) + def loop_fn(i): + qml.Hadamard(wires=i) + + # nested for loops. + # outer for loop updates x + @qml.for_loop(0, n, 1) + def loop_fn_returns(i, x): + qml.RX(x, wires=i) + + # inner for loop + @qml.for_loop(i + 1, n, 1) + def inner(j): + qml.CRY(x**2, [i, j]) + + inner() + + return x + 0.1 + + loop_fn() + loop_fn_returns(x) + + return qml.expval(qml.PauliZ(0)) + + x = 0.5 + assert jnp.allclose(circuit(x, 3), qml.qjit(circuit)(x, 3)) + + res = circuit.tape.operations + expected = [ + qml.Hadamard(wires=[0]), + qml.Hadamard(wires=[1]), + qml.Hadamard(wires=[2]), + qml.RX(0.5, wires=[0]), + qml.CRY(0.25, wires=[0, 1]), + qml.CRY(0.25, wires=[0, 2]), + qml.RX(0.6, wires=[1]), + qml.CRY(0.36, wires=[1, 2]), + qml.RX(0.7, wires=[2]), + ] + + _ = [qml.assert_equal(i, j) for i, j in zip(res, expected)] + def test_cond(self): """Test condition with simple true_fn""" dev = qml.device("lightning.qubit", wires=1) diff --git a/tests/test_debugging.py b/tests/test_debugging.py index 207b97bca04..bdd10a4cfd5 100644 --- a/tests/test_debugging.py +++ b/tests/test_debugging.py @@ -181,7 +181,12 @@ def circuit(): return qml.expval(qml.PauliZ(0)) - qml.snapshots(circuit)(shots=200) + with ( + pytest.warns(UserWarning, match="Requested state or density matrix with finite shots") + if isinstance(dev, qml.devices.default_qutrit.DefaultQutrit) + else nullcontext() + ): + qml.snapshots(circuit)(shots=200) @pytest.mark.parametrize("diff_method", [None, "parameter-shift"]) def test_all_state_measurement_snapshot_pure_qubit_dev(self, dev, diff_method): @@ -275,7 +280,8 @@ def circuit(): @pytest.mark.parametrize("diff_method", [None, "backprop", "parameter-shift", "adjoint"]) def test_default_qubit_legacy_only_supports_state(self, diff_method): - dev = qml.device("default.qubit.legacy", wires=2) + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Use of 'default.qubit"): + dev = qml.device("default.qubit.legacy", wires=2) assert qml.debugging.snapshot._is_snapshot_compatible(dev) @@ -634,12 +640,12 @@ def circuit(): # TODO: fallback to simple `np.allclose` tests once `setRandomSeed` is exposed from the lightning C++ code counts, expvals = tuple(zip(*(qml.snapshots(circuit)().values() for _ in range(50)))) - assert ttest_ind([count["0"] for count in counts], 250).pvalue >= 0.8 - assert ttest_ind(expvals, 0.0).pvalue >= 0.8 + assert ttest_ind([count["0"] for count in counts], 250).pvalue >= 0.75 + assert ttest_ind(expvals, 0.0).pvalue >= 0.75 # Make sure shots are overriden correctly counts, _ = tuple(zip(*(qml.snapshots(circuit)(shots=1000).values() for _ in range(50)))) - assert ttest_ind([count["0"] for count in counts], 500).pvalue >= 0.8 + assert ttest_ind([count["0"] for count in counts], 500).pvalue >= 0.75 @pytest.mark.parametrize("diff_method", ["backprop", "adjoint"]) def test_lightning_qubit_fails_for_state_snapshots_with_adjoint_and_backprop(self, diff_method): diff --git a/tests/test_operation.py b/tests/test_operation.py index 4750b8feb3f..3062803df0c 100644 --- a/tests/test_operation.py +++ b/tests/test_operation.py @@ -321,6 +321,21 @@ class DummyOp(Operator): assert op._ndim_params == (ndim_params,) assert op.ndim_params == (0,) + def test_expand_deprecated(self): + + class MyOp(qml.operation.Operation): + num_wires = 1 + has_decomposition = True + + @staticmethod + def compute_decomposition(*params, wires=None, **hyperparameters): + return [qml.Hadamard(wires=wires)] + + op = MyOp(wires=0) + + with pytest.warns(qml.PennyLaneDeprecationWarning, match="'Operator.expand' is deprecated"): + op.expand() + class TestPytreeMethods: def test_pytree_defaults(self): @@ -1073,7 +1088,7 @@ def test_all_wires_defined_but_init_with_one(self): """Test that an exception is raised if the class is defined with ALL wires, but then instantiated with only one""" - dev1 = qml.device("default.qubit.legacy", wires=2) + dev1 = qml.device("default.qubit", wires=2) class DummyOp(qml.operation.Operation): r"""Dummy custom operator""" diff --git a/tests/test_qnode.py b/tests/test_qnode.py index 966f6d31467..52d9155e448 100644 --- a/tests/test_qnode.py +++ b/tests/test_qnode.py @@ -18,7 +18,6 @@ import warnings from dataclasses import asdict, replace from functools import partial -from typing import Callable, Tuple import numpy as np import pytest @@ -28,7 +27,8 @@ from pennylane import QNode from pennylane import numpy as pnp from pennylane import qnode -from pennylane.tape import QuantumScript +from pennylane.tape import QuantumScript, QuantumTapeBatch +from pennylane.typing import PostprocessingFn from pennylane.workflow.qnode import _prune_dynamic_transform @@ -101,6 +101,15 @@ def f(): assert f.execute_kwargs["cache"] is True + def test_max_expansion_is_deprecated(self): + """Test that a warning is raised when using the deprecated max_expansion argument""" + dev = qml.device("default.qubit", wires=1) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The max_expansion argument is deprecated", + ): + QNode(dummyfunc, dev, max_expansion=10) + # pylint: disable=too-many-public-methods class TestValidation: @@ -878,10 +887,7 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) - @pytest.mark.parametrize( - "dev", - [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)], - ) + @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.legacy"]) @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize( @@ -896,11 +902,12 @@ def circuit(x, y): ], ) def test_defer_meas_if_mcm_unsupported( - self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker + self, dev_name, first_par, sec_par, return_type, mv_return, mv_res, mocker ): # pylint: disable=too-many-arguments """Tests that the transform using the deferred measurement principle is applied if the device doesn't support mid-circuit measurements natively.""" + dev = qml.device(dev_name, wires=3) @qml.qnode(dev) def cry_qnode(x, y): @@ -1388,7 +1395,7 @@ def null_postprocessing(results): @qml.transforms.core.transform def just_pauli_x_out( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing @@ -1418,7 +1425,7 @@ def tet_transform_program_modifies_results(self): @qml.transforms.core.transform def pin_result( tape: qml.tape.QuantumTape, requested_result - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: def postprocessing(_: qml.typing.ResultBatch) -> qml.typing.Result: return requested_result @@ -1444,7 +1451,9 @@ def null_postprocessing(results): return results[0] @qml.transforms.core.transform - def just_pauli_x_out(tape: qml.tape.QuantumTape) -> (Tuple[qml.tape.QuantumTape], Callable): + def just_pauli_x_out( + tape: qml.tape.QuantumTape, + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing @@ -1452,7 +1461,7 @@ def just_pauli_x_out(tape: qml.tape.QuantumTape) -> (Tuple[qml.tape.QuantumTape] @qml.transforms.core.transform def repeat_operations( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: new_tape = qml.tape.QuantumScript( tape.operations + copy.deepcopy(tape.operations), tape.measurements ) @@ -1496,13 +1505,13 @@ def add_shift(results, shift): @qml.transforms.core.transform def scale_output( tape: qml.tape.QuantumTape, factor - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return (tape,), partial(scale_by_factor, factor=factor) @qml.transforms.core.transform def shift_output( tape: qml.tape.QuantumTape, shift - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return (tape,), partial(add_shift, shift=shift) @partial(shift_output, shift=1.0) @@ -1533,7 +1542,7 @@ def num_of_shots_from_sample(results): return len(results[0]) @qml.transforms.core.transform - def use_n_shots(tape: qml.tape.QuantumTape, n) -> (Tuple[qml.tape.QuantumTape], Callable): + def use_n_shots(tape: qml.tape.QuantumTape, n) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript(tape.operations, tape.measurements, shots=n), ), num_of_shots_from_sample @@ -1984,10 +1993,15 @@ def test_device_expansion_strategy(self, mocker): dev = qml.device("default.qubit", wires=2) x = pnp.array(0.5, requires_grad=True) - @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") - def circuit(x): - qml.SingleExcitation(x, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="'expansion_strategy' attribute is deprecated", + ): + + @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") + def circuit(x): + qml.SingleExcitation(x, wires=[0, 1]) + return qml.expval(qml.PauliX(0)) assert circuit.expansion_strategy == "device" assert circuit.execute_kwargs["expand_fn"] is None @@ -2019,10 +2033,15 @@ def transform_program(tapes): dev = qml.device("default.qubit", wires=2) monkeypatch.setattr(dev, "preprocess", preprocess_with_batchtransform) - @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") - def circuit(x): - qml.SingleExcitation(x, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="'expansion_strategy' attribute is deprecated", + ): + + @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") + def circuit(x): + qml.SingleExcitation(x, wires=[0, 1]) + return qml.expval(qml.PauliX(0)) with pytest.raises( ValueError, @@ -2056,14 +2075,14 @@ def test_prune_dynamic_transform(): program1 = qml.transforms.core.TransformProgram( [ qml.transforms.dynamic_one_shot, - qml.transforms.sum_expand, + qml.transforms.split_non_commuting, qml.transforms.dynamic_one_shot, ] ) program2 = qml.transforms.core.TransformProgram( [ qml.transforms.dynamic_one_shot, - qml.transforms.sum_expand, + qml.transforms.split_non_commuting, ] ) @@ -2078,14 +2097,14 @@ def test_prune_dynamic_transform_with_mcm(): program1 = qml.transforms.core.TransformProgram( [ qml.transforms.dynamic_one_shot, - qml.transforms.sum_expand, + qml.transforms.split_non_commuting, qml.devices.preprocess.mid_circuit_measurements, ] ) program2 = qml.transforms.core.TransformProgram( [ qml.transforms.dynamic_one_shot, - qml.transforms.sum_expand, + qml.transforms.split_non_commuting, ] ) diff --git a/tests/test_qnode_legacy.py b/tests/test_qnode_legacy.py index c98fc3f05fd..d8edc0890ee 100644 --- a/tests/test_qnode_legacy.py +++ b/tests/test_qnode_legacy.py @@ -17,7 +17,6 @@ # pylint: disable=import-outside-toplevel, protected-access, no-member import warnings from functools import partial -from typing import Callable, Tuple import numpy as np import pytest @@ -28,7 +27,8 @@ from pennylane import numpy as pnp from pennylane import qnode from pennylane.resource import Resources -from pennylane.tape import QuantumScript +from pennylane.tape import QuantumScript, QuantumTapeBatch +from pennylane.typing import PostprocessingFn def dummyfunc(): @@ -491,13 +491,14 @@ def test_adjoint_finite_shots(self): dev = qml.device("default.qubit.legacy", wires=1, shots=1) - @qnode(dev, diff_method="adjoint") - def circ(): - return qml.expval(qml.PauliZ(0)) - with pytest.warns( UserWarning, match="Requested adjoint differentiation to be computed with finite shots." ): + + @qnode(dev, diff_method="adjoint") + def circ(): + return qml.expval(qml.PauliZ(0)) + circ() @pytest.mark.autograd @@ -621,7 +622,8 @@ def circuit(params): } def test_autograd_interface_device_switched_no_warnings(self): - """Test that checks that no warning is raised for device switch when you define an interface.""" + """Test that checks that no warning is raised for device switch when you define an interface, + except for the deprecation warnings which will be caught by the fixture.""" dev = qml.device("default.qubit.legacy", wires=1) @qml.qnode(dev, interface="autograd") @@ -629,17 +631,16 @@ def circuit(params): qml.RX(params, wires=0) return qml.expval(qml.PauliZ(0)) - with warnings.catch_warnings(record=True) as record: - circuit(qml.numpy.array(0.1, requires_grad=True)) - - assert len(record) == 0 + circuit(qml.numpy.array(0.1, requires_grad=True)) def test_not_giving_mode_kwarg_does_not_raise_warning(self): - """Test that not providing a value for mode does not raise a warning.""" + """Test that not providing a value for mode does not raise a warning + except for the deprecation warning.""" with warnings.catch_warnings(record=True) as record: - _ = qml.QNode(lambda f: f, qml.device("default.qubit.legacy", wires=1)) + qml.QNode(lambda f: f, qml.device("default.qubit.legacy", wires=1)) - assert len(record) == 0 + assert len(record) == 1 + assert record[0].category == qml.PennyLaneDeprecationWarning class TestTapeConstruction: @@ -1089,9 +1090,7 @@ def circuit(): assert len(circuit.tape.operations) == 2 assert isinstance(circuit.tape.operations[1], qml.measurements.MidMeasureMP) - @pytest.mark.parametrize( - "dev", [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)] - ) + @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.legacy"]) @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize( @@ -1106,12 +1105,14 @@ def circuit(): ], ) def test_defer_meas_if_mcm_unsupported( - self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker + self, dev_name, first_par, sec_par, return_type, mv_return, mv_res, mocker ): # pylint: disable=too-many-arguments """Tests that the transform using the deferred measurement principle is applied if the device doesn't support mid-circuit measurements natively.""" + dev = qml.device(dev_name, wires=3) + @qml.qnode(dev) def cry_qnode(x, y): """QNode where we apply a controlled Y-rotation.""" @@ -1548,7 +1549,7 @@ def null_postprocessing(results): @qml.transforms.core.transform def just_pauli_x_out( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing @@ -1578,7 +1579,7 @@ def tet_transform_program_modifies_results(self): @qml.transforms.core.transform def pin_result( tape: qml.tape.QuantumTape, requested_result - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: def postprocessing(_: qml.typing.ResultBatch) -> qml.typing.Result: return requested_result @@ -1604,7 +1605,9 @@ def null_postprocessing(results): return results[0] @qml.transforms.core.transform - def just_pauli_x_out(tape: qml.tape.QuantumTape) -> (Tuple[qml.tape.QuantumTape], Callable): + def just_pauli_x_out( + tape: qml.tape.QuantumTape, + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript([qml.PauliX(0)], tape.measurements), ), null_postprocessing @@ -1612,7 +1615,7 @@ def just_pauli_x_out(tape: qml.tape.QuantumTape) -> (Tuple[qml.tape.QuantumTape] @qml.transforms.core.transform def repeat_operations( tape: qml.tape.QuantumTape, - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: new_tape = qml.tape.QuantumScript( tape.operations + copy.deepcopy(tape.operations), tape.measurements ) @@ -1656,13 +1659,13 @@ def add_shift(results, shift): @qml.transforms.core.transform def scale_output( tape: qml.tape.QuantumTape, factor - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return (tape,), partial(scale_by_factor, factor=factor) @qml.transforms.core.transform def shift_output( tape: qml.tape.QuantumTape, shift - ) -> (Tuple[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: return (tape,), partial(add_shift, shift=shift) @partial(shift_output, shift=1.0) @@ -1693,7 +1696,7 @@ def num_of_shots_from_sample(results): return len(results[0]) @qml.transforms.core.transform - def use_n_shots(tape: qml.tape.QuantumTape, n) -> (Tuple[qml.tape.QuantumTape], Callable): + def use_n_shots(tape: qml.tape.QuantumTape, n) -> tuple[QuantumTapeBatch, PostprocessingFn]: return ( qml.tape.QuantumScript(tape.operations, tape.measurements, shots=n), ), num_of_shots_from_sample @@ -1888,10 +1891,15 @@ def test_device_expansion_strategy(self, mocker): dev = qml.device("default.qubit.legacy", wires=2) x = pnp.array(0.5, requires_grad=True) - @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") - def circuit(x): - qml.SingleExcitation(x, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="'expansion_strategy' attribute is deprecated", + ): + + @qnode(dev, diff_method="parameter-shift", expansion_strategy="device") + def circuit(x): + qml.SingleExcitation(x, wires=[0, 1]) + return qml.expval(qml.PauliX(0)) assert circuit.expansion_strategy == "device" assert circuit.execute_kwargs["expand_fn"] is None diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py index 2b2387ff6ed..787ed809fbf 100644 --- a/tests/test_qubit_device.py +++ b/tests/test_qubit_device.py @@ -1211,7 +1211,7 @@ def test_device_executions(self): """Test the number of times a qubit device is executed over a QNode's lifetime is tracked by `num_executions`""" - dev_1 = qml.device("default.qubit.legacy", wires=2) + dev_1 = qml.device("default.mixed", wires=2) def circuit_1(x, y): qml.RX(x, wires=[0]) @@ -1227,7 +1227,7 @@ def circuit_1(x, y): assert dev_1.num_executions == num_evals_1 # test a second instance of a default qubit device - dev_2 = qml.device("default.qubit.legacy", wires=2) + dev_2 = qml.device("default.mixed", wires=2) def circuit_2(x): qml.RX(x, wires=[0]) @@ -1272,7 +1272,7 @@ def test_device_executions(self): """Test the number of times a qubit device is executed over a QNode's lifetime is tracked by `num_executions`""" - dev_1 = qml.device("default.qubit.legacy", wires=2) + dev_1 = qml.device("default.mixed", wires=2) def circuit_1(x, y): qml.RX(x, wires=[0]) @@ -1285,10 +1285,10 @@ def circuit_1(x, y): for _ in range(num_evals_1): node_1(0.432, np.array([0.12, 0.5, 3.2])) - assert dev_1.num_executions == num_evals_1 + assert dev_1.num_executions == num_evals_1 * 3 # test a second instance of a default qubit device - dev_2 = qml.device("default.qubit.legacy", wires=2) + dev_2 = qml.device("default.mixed", wires=2) assert dev_2.num_executions == 0 @@ -1302,7 +1302,7 @@ def circuit_2(x, y): for _ in range(num_evals_2): node_2(np.array([0.432, 0.61, 8.2]), 0.12) - assert dev_2.num_executions == num_evals_2 + assert dev_2.num_executions == num_evals_2 * 3 # test a new circuit on an existing instance of a qubit device def circuit_3(x, y): @@ -1315,7 +1315,7 @@ def circuit_3(x, y): for _ in range(num_evals_3): node_3(np.array([0.432, 0.2]), np.array([0.12, 1.214])) - assert dev_1.num_executions == num_evals_1 + num_evals_3 + assert dev_1.num_executions == num_evals_1 * 3 + num_evals_3 * 2 class TestBatchExecution: @@ -1504,7 +1504,7 @@ def test_tracker_multi_execution(self, dev_name): @pytest.mark.autograd def test_tracker_grad(self): """Test that the tracker can track resources through a gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1, shots=100) + dev = qml.device("default.qubit", wires=1, shots=100) @qml.qnode(dev, diff_method="parameter-shift") def circuit(x): @@ -1540,8 +1540,9 @@ def test_samples_to_counts_with_nan(self): """Test that the counts function disregards failed measurements (samples including NaN values) when totalling counts""" # generate 1000 samples for 2 wires, randomly distributed between 0 and 1 - device = qml.device("default.qubit.legacy", wires=2, shots=1000) - device._state = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j] + device = qml.device("default.mixed", wires=2, shots=1000) + sv = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j] + device._state = np.outer(sv, sv) device._samples = device.generate_samples() samples = device.sample(qml.measurements.CountsMP()) @@ -1568,9 +1569,12 @@ def test_samples_to_counts_with_many_wires(self, all_outcomes): # generate 1000 samples for 10 wires, randomly distributed between 0 and 1 n_wires = 10 shots = 100 - device = qml.device("default.qubit.legacy", wires=n_wires, shots=shots) - state = np.random.rand(*([2] * n_wires)) - device._state = state / np.linalg.norm(state) + device = qml.device("default.mixed", wires=n_wires, shots=shots) + + sv = np.random.rand(*([2] * n_wires)) + state = sv / np.linalg.norm(sv) + + device._state = np.outer(state, state) device._samples = device.generate_samples() samples = device.sample(qml.measurements.CountsMP(all_outcomes=all_outcomes)) diff --git a/tests/transforms/core/test_transform_dispatcher.py b/tests/transforms/core/test_transform_dispatcher.py index 94f7d74e7f5..a4f23db6cfd 100644 --- a/tests/transforms/core/test_transform_dispatcher.py +++ b/tests/transforms/core/test_transform_dispatcher.py @@ -13,14 +13,15 @@ # limitations under the License. """Unit and integration tests for the transform dispatcher.""" import inspect +from collections.abc import Callable, Sequence from functools import partial -from typing import Callable, Sequence import pytest import pennylane as qml +from pennylane.tape import QuantumTapeBatch from pennylane.transforms.core import TransformContainer, TransformError, transform -from pennylane.typing import TensorLike +from pennylane.typing import PostprocessingFn, TensorLike dev = qml.device("default.qubit", wires=2) @@ -49,7 +50,7 @@ def qfunc_circuit(a: qml.typing.TensorLike): def no_tape_transform( circuit: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Transform without tape.""" circuit = circuit.copy() circuit._ops.pop(index) # pylint:disable=protected-access @@ -58,27 +59,29 @@ def no_tape_transform( def no_quantum_tape_transform( tape: qml.operation.Operator, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, Callable]: """Transform with wrong hinting.""" tape = tape.copy() tape._ops.pop(index) # pylint:disable=protected-access return [tape], lambda x: x -def no_processing_fn_transform(tape: qml.tape.QuantumTape) -> Sequence[qml.tape.QuantumTape]: +def no_processing_fn_transform(tape: qml.tape.QuantumTape) -> QuantumTapeBatch: """Transform without processing fn.""" tape_copy = tape.copy() return [tape, tape_copy] -def no_tape_sequence_transform(tape: qml.tape.QuantumTape) -> (qml.tape.QuantumTape, Callable): +def no_tape_sequence_transform( + tape: qml.tape.QuantumTape, +) -> tuple[qml.tape.QuantumTape, PostprocessingFn]: """Transform wihtout Sequence return.""" return tape, lambda x: x def no_callable_return( tape: qml.tape.QuantumTape, -) -> (Sequence[qml.tape.QuantumTape], qml.tape.QuantumTape): +) -> tuple[QuantumTapeBatch, qml.tape.QuantumTape]: """Transform without callable return.""" return list(tape), tape @@ -99,7 +102,7 @@ def no_callable_return( def first_valid_transform( tape: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid transform.""" tape = tape.copy() tape._ops.pop(index) # pylint:disable=protected-access @@ -109,7 +112,7 @@ def first_valid_transform( def second_valid_transform( tape: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid trasnform.""" tape1 = tape.copy() tape2 = tape.copy() @@ -128,7 +131,7 @@ def fn(results): # Valid expand transform def expand_transform( tape: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """Multiple args expand fn.""" tape._ops.pop(index) # pylint:disable=protected-access return [tape], lambda x: x @@ -137,14 +140,14 @@ def expand_transform( # Non-valid expand transform def non_valid_expand_transform( tape: qml.tape.QuantumTape, -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid expand transform.""" return [tape], lambda x: x ########################################## # Valid informative transform -def informative_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable): +def informative_transform(tape: qml.tape.QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid informative transform""" def fn(results): @@ -542,7 +545,7 @@ def comb_postproc(results: TensorLike, fn1: Callable, fn2: Callable): tape = qml.tape.QuantumTape(ops, measur) - batch1, fn1 = qml.transforms.hamiltonian_expand(tape) + batch1, fn1 = qml.transforms.split_non_commuting(tape) assert check_batch(batch1) batch2, fn2 = qml.transforms.merge_rotations(batch1) @@ -717,7 +720,7 @@ def test_sphinx_build(self, monkeypatch): @qml.transforms.core.transform def custom_transform( # pylint:disable=unused-variable tape: qml.tape.QuantumTape, index: int - ) -> (Sequence[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid transform.""" tape = tape.copy() tape._ops.pop(index) # pylint:disable=protected-access diff --git a/tests/transforms/core/test_transform_program.py b/tests/transforms/core/test_transform_program.py index b104213f845..057d4fa8859 100644 --- a/tests/transforms/core/test_transform_program.py +++ b/tests/transforms/core/test_transform_program.py @@ -13,12 +13,11 @@ # limitations under the License. """Unit and integration tests for the transform program.""" # pylint: disable=no-member -from typing import Callable, Sequence import pytest import pennylane as qml -from pennylane.tape import QuantumScript +from pennylane.tape import QuantumScript, QuantumTapeBatch from pennylane.transforms.core import ( TransformContainer, TransformError, @@ -30,12 +29,12 @@ _batch_postprocessing, null_postprocessing, ) -from pennylane.typing import Result, ResultBatch +from pennylane.typing import PostprocessingFn, Result, ResultBatch def first_valid_transform( tape: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid transform.""" tape = tape.copy() tape._ops.pop(index) # pylint:disable=protected-access @@ -44,14 +43,14 @@ def first_valid_transform( def expand_transform( tape: qml.tape.QuantumTape, index: int # pylint:disable=unused-argument -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid expand transform.""" return [tape], lambda x: x def second_valid_transform( tape: qml.tape.QuantumTape, index: int -) -> (Sequence[qml.tape.QuantumTape], Callable): +) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid trasnform.""" tape1 = tape.copy() tape2 = tape.copy() @@ -63,7 +62,7 @@ def fn(results): return [tape1, tape2], fn -def informative_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable): +def informative_transform(tape: qml.tape.QuantumTape) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid informative transform""" def fn(results): @@ -595,7 +594,7 @@ def single_null_postprocessing(results): def remove_operation_at_index( tape: qml.tape.QuantumTape, index: int - ) -> (Sequence[qml.tape.QuantumTape], Callable): + ) -> tuple[QuantumTapeBatch, PostprocessingFn]: """A valid transform.""" new_ops = list(tape.operations) new_ops.pop(index) # pylint:disable=protected-access diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py index 0fe8739b6ca..b3f2ebad99d 100644 --- a/tests/transforms/test_hamiltonian_expand.py +++ b/tests/transforms/test_hamiltonian_expand.py @@ -15,6 +15,7 @@ Unit tests for the ``hamiltonian_expand`` transform. """ import functools +import warnings import numpy as np import pytest @@ -90,6 +91,18 @@ class TestHamiltonianExpand: """Tests for the hamiltonian_expand transform""" + @pytest.fixture(scope="function", autouse=True) + def capture_warnings(self): + with pytest.warns(qml.PennyLaneDeprecationWarning) as record: + yield + + for w in record: + assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + if "qml.transforms.hamiltonian_expand is deprecated" not in str(w.message): + warnings.warn(w.message, w.category) + else: + assert "qml.transforms.hamiltonian_expand is deprecated" in str(w.message) + def test_ham_with_no_terms_raises(self): """Tests that the hamiltonian_expand transform raises an error for a Hamiltonian with no terms.""" mps = [qml.expval(qml.Hamiltonian([], []))] @@ -518,6 +531,18 @@ def circuit(): class TestSumExpand: """Tests for the sum_expand transform""" + @pytest.fixture(scope="function", autouse=True) + def capture_warnings(self): + with pytest.warns(qml.PennyLaneDeprecationWarning) as record: + yield + + for w in record: + assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + if "qml.transforms.sum_expand is deprecated" not in str(w.message): + warnings.warn(w.message, w.category) + else: + assert "qml.transforms.sum_expand is deprecated" in str(w.message) + def test_observables_on_same_wires(self): """Test that even if the observables are on the same wires, if they are different operations, they are separated. This is testing for a case that gave rise to a bug that occured due to a problem in MeasurementProcess.hash. @@ -544,6 +569,7 @@ def test_sums(self, qscript, output): assert all(qml.math.allclose(o, e) for o, e in zip(output, expval)) @pytest.mark.parametrize(("qscript", "output"), zip(SUM_QSCRIPTS, SUM_OUTPUTS)) + @pytest.mark.filterwarnings("ignore:Use of 'default.qubit.legacy' is deprecated") def test_sums_legacy_opmath(self, qscript, output): """Tests that the sum_expand transform returns the correct value""" dev_old = qml.device("default.qubit.legacy", wires=4) diff --git a/tests/transforms/test_mitigate.py b/tests/transforms/test_mitigate.py index 7cf6dd17915..002a7b1af06 100644 --- a/tests/transforms/test_mitigate.py +++ b/tests/transforms/test_mitigate.py @@ -23,7 +23,12 @@ import pennylane as qml from pennylane import numpy as np from pennylane.tape import QuantumScript -from pennylane.transforms import fold_global, mitigate_with_zne, richardson_extrapolate +from pennylane.transforms import ( + exponential_extrapolate, + fold_global, + mitigate_with_zne, + richardson_extrapolate, +) with qml.queuing.AnnotatedQueue() as q_tape: qml.BasisState([1], wires=0) @@ -109,7 +114,8 @@ def test_extrapolate_call(self, mocker): for t in tapes: same_tape(t, tape) - def test_multi_returns(self): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_multi_returns(self, extrapolate): """Tests if the expected shape is returned when mitigating a circuit with two returns""" noise_strength = 0.05 @@ -126,7 +132,7 @@ def test_multi_returns(self): qml.transforms.mitigate_with_zne, scale_factors=[1, 2, 3], folding=fold_global, - extrapolate=richardson_extrapolate, + extrapolate=extrapolate, ) @qml.qnode(dev) def mitigated_circuit(w1, w2): @@ -507,17 +513,67 @@ def test_polyfit(self): coeffs = qml.transforms.mitigate._polyfit(x, y, 2) assert qml.math.allclose(qml.math.squeeze(coeffs), [3, 2, 1]) + @pytest.mark.parametrize("exp_params", [[0.5, -2, 2], [-9, -4, 0]]) + def test_exponential_extrapolation_accuracy(self, exp_params): + """Testing the exponential extrapolation works as expected for known exponential models.""" + A, B, asymptote = exp_params + x = np.linspace(1, 4, 4) + y = A * np.exp(B * x) + asymptote + zne_val = qml.transforms.exponential_extrapolate(x, y, asymptote=asymptote) + assert qml.math.allclose(zne_val, A + asymptote, atol=1e-3) + + @pytest.mark.autograd + def test_exponential_extrapolation_autograd(self): + """Test exponential extrapolation works with expvals stored as a numpy array.""" + scale_factors = [1, 3, 5] + noise_scaled_expvals = np.array([0.9, 0.8, 0.7]) + zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals) + assert isinstance(zne_val, np.ndarray) + assert zne_val.ndim == 0 + + @pytest.mark.tf + def test_exponential_extrapolation_tf(self): + """Test exponential extrapolation works with expvals stored as a tensorflow tensor.""" + import tensorflow as tf + + scale_factors = [1, 3, 5] + noise_scaled_expvals = tf.constant([0.9, 0.8, 0.7], dtype=tf.float32) + zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals) + assert tf.is_tensor(zne_val) + assert zne_val.shape.ndims == 0 + + @pytest.mark.torch + def test_exponential_extrapolation_torch(self): + """Test exponential extrapolation works with expvals stored as a torch tensor.""" + import torch + + scale_factors = [1, 3, 5] + noise_scaled_expvals = torch.tensor([0.9, 0.8, 0.7]) + zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals) + assert torch.is_tensor(zne_val) + assert zne_val.ndimension() == 0 + + @pytest.mark.jax + def test_exponential_extrapolation_jax(self): + """Test exponential extrapolation works with expvals stored as a jax array.""" + import jax.numpy as jnp + + scale_factors = [1, 3, 5] + noise_scaled_expvals = jnp.array([0.9, 0.8, 0.7]) + zne_val = qml.transforms.exponential_extrapolate(scale_factors, noise_scaled_expvals) + assert isinstance(zne_val, jnp.ndarray) + assert zne_val.ndim == 0 + @pytest.mark.autograd - def test_diffability_autograd(self): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_autograd(self, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in autograd""" qnode_noisy = qml.QNode(qfunc, dev_noisy) qnode_ideal = qml.QNode(qfunc, dev_ideal) scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = np.array([np.pi / 4, np.pi / 4], requires_grad=True) @@ -531,7 +587,8 @@ def test_diffability_autograd(self): @pytest.mark.jax @pytest.mark.parametrize("interface", ["auto", "jax"]) - def test_diffability_jax(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_jax(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax""" import jax import jax.numpy as jnp @@ -541,9 +598,7 @@ def test_diffability_jax(self, interface): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = jnp.array( [np.pi / 4, np.pi / 4], @@ -559,7 +614,8 @@ def test_diffability_jax(self, interface): @pytest.mark.jax @pytest.mark.parametrize("interface", ["auto", "jax", "jax-jit"]) - def test_diffability_jaxjit(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_jaxjit(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax-jit""" import jax import jax.numpy as jnp @@ -570,7 +626,7 @@ def test_diffability_jaxjit(self, interface): scale_factors = [1.0, 2.0, 3.0] mitigated_qnode = jax.jit( - mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate) + mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) ) theta = jnp.array( @@ -587,7 +643,8 @@ def test_diffability_jaxjit(self, interface): @pytest.mark.torch @pytest.mark.parametrize("interface", ["auto", "torch"]) - def test_diffability_torch(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_torch(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in torch""" import torch @@ -596,9 +653,7 @@ def test_diffability_torch(self, interface): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = torch.tensor([np.pi / 4, np.pi / 4], requires_grad=True) @@ -616,7 +671,8 @@ def test_diffability_torch(self, interface): @pytest.mark.tf @pytest.mark.parametrize("interface", ["auto", "tf"]) - def test_diffability_tf(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_tf(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in tf""" import tensorflow as tf @@ -625,9 +681,7 @@ def test_diffability_tf(self, interface): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = tf.Variable([np.pi / 4, np.pi / 4]) @@ -645,7 +699,8 @@ def test_diffability_tf(self, interface): assert qml.math.allclose(grad, grad_ideal, atol=1e-2) @pytest.mark.autograd - def test_diffability_autograd_multi(self): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_autograd_multi(self, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in autograd for multiple measurements""" qnode_noisy = qml.QNode(qfunc_multi, dev_noisy) @@ -653,9 +708,7 @@ def test_diffability_autograd_multi(self): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = np.array([np.pi / 4, np.pi / 6], requires_grad=True) @@ -669,7 +722,8 @@ def test_diffability_autograd_multi(self): @pytest.mark.jax @pytest.mark.parametrize("interface", ["auto", "jax"]) - def test_diffability_jax_multi(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_jax_multi(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax for multiple measurements""" import jax @@ -680,9 +734,7 @@ def test_diffability_jax_multi(self, interface): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = jnp.array( [np.pi / 4, np.pi / 6], @@ -698,7 +750,8 @@ def test_diffability_jax_multi(self, interface): @pytest.mark.jax @pytest.mark.parametrize("interface", ["auto", "jax", "jax-jit"]) - def test_diffability_jaxjit_multi(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_jaxjit_multi(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in jax-jit for multiple measurements""" import jax @@ -710,7 +763,7 @@ def test_diffability_jaxjit_multi(self, interface): scale_factors = [1.0, 2.0, 3.0] mitigated_qnode = jax.jit( - mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate) + mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) ) theta = jnp.array( @@ -727,7 +780,8 @@ def test_diffability_jaxjit_multi(self, interface): @pytest.mark.torch @pytest.mark.parametrize("interface", ["auto", "torch"]) - def test_diffability_torch_multi(self, interface): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_torch_multi(self, interface, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in torch for multiple measurements""" import torch @@ -737,9 +791,7 @@ def test_diffability_torch_multi(self, interface): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = torch.tensor([np.pi / 4, np.pi / 6], requires_grad=True) @@ -756,7 +808,8 @@ def test_diffability_torch_multi(self, interface): assert qml.math.allclose(grad, grad_ideal, atol=1e-2) @pytest.mark.tf - def test_diffability_tf_multi(self): + @pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate]) + def test_diffability_tf_multi(self, extrapolate): """Testing that the mitigated qnode can be differentiated and returns the correct gradient in tf for multiple measurements""" import tensorflow as tf @@ -766,9 +819,7 @@ def test_diffability_tf_multi(self): scale_factors = [1.0, 2.0, 3.0] - mitigated_qnode = mitigate_with_zne( - qnode_noisy, scale_factors, fold_global, richardson_extrapolate - ) + mitigated_qnode = mitigate_with_zne(qnode_noisy, scale_factors, fold_global, extrapolate) theta = tf.Variable([np.pi / 4, np.pi / 6]) diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py index 6439bb04e2d..1a78c1ab4e2 100644 --- a/tests/transforms/test_qcut.py +++ b/tests/transforms/test_qcut.py @@ -4070,9 +4070,7 @@ def test_simple_cut_circuit_torch_trace(self, mocker, use_opt_einsum): import torch - # TODO: this passes with default.qubit locally, but fails on CI - # possibly an architecture-specific issue - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) @qml.qnode(dev, interface="torch") def circuit(x): @@ -5883,7 +5881,7 @@ def block(weights, wires): tape0 = qml.tape.QuantumScript.from_queue(q0) tape = tape0.expand() - tapes, _ = qml.transforms.hamiltonian_expand(tape, group=False) + tapes, _ = qml.transforms.split_non_commuting(tape, grouping_strategy=None) frag_lens = [5, 7] frag_ords = [[1, 6], [3, 6]] diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py index 00393889825..4bf16a061d4 100644 --- a/tests/transforms/test_tape_expand.py +++ b/tests/transforms/test_tape_expand.py @@ -17,6 +17,8 @@ # pylint: disable=too-few-public-methods, invalid-unary-operand-type, no-member, # pylint: disable=arguments-differ, arguments-renamed, +import warnings + import numpy as np import pytest @@ -429,6 +431,21 @@ def custom_basic_entangler_layers(weights, wires, **kwargs): class TestCreateCustomDecompExpandFn: """Tests for the custom_decomps argument for devices""" + @pytest.fixture(scope="function", autouse=True) + def capture_warnings(self): + with pytest.warns( + qml.PennyLaneDeprecationWarning, + ) as record: + yield + + assert any( + "'expansion_strategy' attribute is deprecated" in str(w.message) for w in record + ) + + for w in record: + if "'expansion_strategy' attribute is deprecated" not in str(w.message): + warnings.warn(w.message, w.category) + @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"]) def test_string_and_operator_allowed(self, device_name): """Test that the custom_decomps dictionary accepts both strings and operator classes as keys.""" @@ -535,7 +552,12 @@ def test_no_decomp_with_depth_zero(self, device_name): """Test that specifying a single custom decomposition works as expected.""" custom_decomps = {"Hadamard": custom_hadamard, "CNOT": custom_cnot} - decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=0) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The decomp_depth argument is deprecated" + ): + decomp_dev = qml.device( + device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=0 + ) @qml.qnode(decomp_dev, expansion_strategy="device") def circuit(): @@ -712,13 +734,16 @@ def test_custom_decomp_different_depth(self, device_name): # not be further decomposed even though the custom decomposition is specified. custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx} - decomp_dev_2 = qml.device( - device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=2 - ) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The decomp_depth argument is deprecated" + ): + decomp_dev_2 = qml.device( + device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=2 + ) - decomp_dev_3 = qml.device( - device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=3 - ) + decomp_dev_3 = qml.device( + device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=3 + ) def circuit(): qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1]) @@ -901,7 +926,7 @@ def test_custom_decomp_used_twice(self): custom_decomps = {"MultiRZ": qml.MultiRZ.compute_decomposition} dev = qml.device("lightning.qubit", wires=2, custom_decomps=custom_decomps) - @qml.qnode(dev, diff_method="adjoint") + @qml.qnode(dev, diff_method="adjoint", expansion_strategy="gradient") def cost(theta): qml.Hadamard(wires=0) qml.Hadamard(wires=1) diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py index 186c362e705..fd782104df4 100644 --- a/tests/workflow/test_cache_transform.py +++ b/tests/workflow/test_cache_transform.py @@ -15,7 +15,7 @@ Unit tests for the :func:`_cache_transform` and :func:`_apply_cache_transform` functions. """ # pylint: disable=protected-access,redefined-outer-name -from typing import MutableMapping +from collections.abc import MutableMapping from unittest.mock import MagicMock import pytest diff --git a/tests/workflow/test_construct_batch.py b/tests/workflow/test_construct_batch.py index 8e8400b2c52..daef883b333 100644 --- a/tests/workflow/test_construct_batch.py +++ b/tests/workflow/test_construct_batch.py @@ -133,7 +133,7 @@ def test_get_transform_program_device_gradient(self): dev = qml.device("default.qubit") - @qml.transforms.sum_expand + @qml.transforms.split_non_commuting @qml.qnode(dev, diff_method="adjoint", device_vjp=False) def circuit(x): qml.RX(x, 0) @@ -150,7 +150,7 @@ def circuit(x): dev_program = dev.preprocess(config)[0] expected = TransformProgram() - expected.add_transform(qml.transforms.sum_expand) + expected.add_transform(qml.transforms.split_non_commuting) expected += dev_program assert full_prog == expected