Source code for pennylane.math.grad
# Copyright 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 submodule defines grad and jacobian for differentiating circuits in an interface-independent way.
"""
from typing import Callable, Sequence, Union
from pennylane._grad import grad as _autograd_grad
from pennylane._grad import jacobian as _autograd_jacobian
from .interface_utils import get_interface
# pylint: disable=import-outside-toplevel
[docs]def grad(f: Callable, argnums: Union[Sequence[int], int] = 0) -> Callable:
"""Compute the gradient in a jax-like manner for any interface.
Args:
f (Callable): a function with a single 0-D scalar output
argnums (Sequence[int] | int ) = 0 : which arguments to differentiate
Returns:
Callable: a function with the same signature as ``f`` that returns the gradient.
.. seealso:: :func:`pennylane.math.jacobian`
Note that this function follows the same design as jax. By default, the function will return the gradient
of the first argument, whether or not other arguments are trainable.
>>> import jax, torch, tensorflow as tf
>>> def f(x, y):
... return x * y
>>> qml.math.grad(f)(qml.numpy.array(2.0), qml.numpy.array(3.0))
tensor(3., requires_grad=True)
>>> qml.math.grad(f)(jax.numpy.array(2.0), jax.numpy.array(3.0))
Array(3., dtype=float32, weak_type=True)
>>> qml.math.grad(f)(torch.tensor(2.0, requires_grad=True), torch.tensor(3.0, requires_grad=True))
tensor(3.)
>>> qml.math.grad(f)(tf.Variable(2.0), tf.Variable(3.0))
<tf.Tensor: shape=(), dtype=float32, numpy=3.0>
``argnums`` can be provided to differentiate multiple arguments.
>>> qml.math.grad(f, argnums=(0,1))(torch.tensor(2.0, requires_grad=True), torch.tensor(3.0, requires_grad=True))
(tensor(3.), tensor(2.))
Note that the selected arguments *must* be of an appropriately trainable datatype, or an error may occur.
>>> qml.math.grad(f)(torch.tensor(1.0), torch.tensor(2.))
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
"""
argnums_integer = False
if isinstance(argnums, int):
argnums = (argnums,)
argnums_integer = True
def compute_grad(*args, **kwargs):
interface = get_interface(*args)
if interface == "autograd":
g = _autograd_grad(f, argnum=argnums)(*args, **kwargs)
return g[0] if argnums_integer else g
if interface == "jax":
import jax
g = jax.grad(f, argnums=argnums)(*args, **kwargs)
return g[0] if argnums_integer else g
if interface == "torch":
y = f(*args, **kwargs)
y.backward()
g = tuple(args[i].grad for i in argnums)
return g[0] if argnums_integer else g
if interface == "tensorflow":
import tensorflow as tf
with tf.GradientTape() as tape:
y = f(*args, **kwargs)
g = tape.gradient(y, tuple(args[i] for i in argnums))
return g[0] if argnums_integer else g
raise ValueError(f"Interface {interface} is not differentiable.")
return compute_grad
# pylint: disable=import-outside-toplevel
def _torch_jac(f, argnums, args, kwargs):
"""Calculate a jacobian via torch."""
from torch.autograd.functional import jacobian as _torch_jac
argnums_torch = (argnums,) if isinstance(argnums, int) else argnums
trainable_args = tuple(args[i] for i in argnums_torch)
# keep track of output type to know how to unpack
output_type_cache = []
def partial_f(*_trainables):
full_args = list(args)
for argnum, value in zip(argnums_torch, _trainables, strict=True):
full_args[argnum] = value
result = f(*full_args, **kwargs)
output_type_cache.append(type(result))
return result
jac = _torch_jac(partial_f, trainable_args)
if output_type_cache[-1] is tuple:
return tuple(j[0] for j in jac) if isinstance(argnums, int) else jac
# else array
return jac[0] if isinstance(argnums, int) else jac
# pylint: disable=import-outside-toplevel
def _tensorflow_jac(f, argnums, args, kwargs):
"""Calculate a jacobian via tensorflow"""
import tensorflow as tf
with tf.GradientTape() as tape:
y = f(*args, **kwargs)
if get_interface(y) != "tensorflow":
raise ValueError(
f"qml.math.jacobian does not work with tensorflow and non-tensor outputs. Got {y} of type {type(y)}."
)
argnums_integer = False
if isinstance(argnums, int):
argnums_tf = (argnums,)
argnums_integer = True
else:
argnums_tf = argnums
g = tape.jacobian(y, tuple(args[i] for i in argnums_tf))
return g[0] if argnums_integer else g
# pylint: disable=import-outside-toplevel
[docs]def jacobian(f: Callable, argnums: Union[Sequence[int], int] = 0) -> Callable:
"""Compute the Jacobian in a jax-like manner for any interface.
Args:
f (Callable): a function with a vector valued output
argnums (Sequence[int] | int ) = 0 : which arguments to differentiate
Returns:
Callable: a function with the same signature as ``f`` that returns the jacobian
.. seealso:: :func:`pennylane.math.grad`
Note that this function follows the same design as jax. By default, the function will return the gradient
of the first argument, whether or not other arguments are trainable.
>>> import jax, torch, tensorflow as tf
>>> def f(x, y):
... return x * y
>>> qml.math.jacobian(f)(qml.numpy.array([2.0, 3.0]), qml.numpy.array(3.0))
array([[3., 0.],
[0., 3.]])
>>> qml.math.jacobian(f)(jax.numpy.array([2.0, 3.0]), jax.numpy.array(3.0))
Array([[3., 0.],
[0., 3.]], dtype=float32)
>>> x_torch = torch.tensor([2.0, 3.0], requires_grad=True)
>>> y_torch = torch.tensor(3.0, requires_grad=True)
>>> qml.math.jacobian(f)(x_torch, y_torch)
tensor([[3., 0.],
[0., 3.]])
>>> qml.math.jacobian(f)(tf.Variable([2.0, 3.0]), tf.Variable(3.0))
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[3., 0.],
[0., 3.]], dtype=float32)>
``argnums`` can be provided to differentiate multiple arguments.
>>> qml.math.jacobian(f, argnums=(0,1))(x_torch, y_torch)
(tensor([[3., 0.],
[0., 3.]]),
tensor([2., 3.]))
While jax can handle taking jacobians of outputs with any pytree shape:
>>> def pytree_f(x):
... return {"a": 2*x, "b": 3*x}
>>> qml.math.jacobian(pytree_f)(jax.numpy.array(2.0))
{'a': Array(2., dtype=float32, weak_type=True),
'b': Array(3., dtype=float32, weak_type=True)}
Torch can only differentiate arrays and tuples:
>>> def tuple_f(x):
... return x**2, x**3
>>> qml.math.jacobian(tuple_f)(torch.tensor(2.0))
(tensor(4.), tensor(12.))
>>> qml.math.jacobian(pytree_f)(torch.tensor(2.0))
TypeError: The outputs of the user-provided function given to jacobian must be
either a Tensor or a tuple of Tensors but the given outputs of the user-provided
function has type <class 'dict'>.
But tensorflow and autograd can only handle array-valued outputs:
>>> qml.math.jacobian(tuple_f)(qml.numpy.array(2.0))
ValueError: autograd can only differentiate with respect to arrays, not <class 'tuple'>
>>> qml.math.jacobian(tuple_f)(tf.Variable(2.0))
ValueError: qml.math.jacobian does not work with tensorflow and non-tensor outputs.
Got (<tf.Tensor: shape=(), dtype=float32, numpy=4.0>,
<tf.Tensor: shape=(), dtype=float32, numpy=8.0>) of type <class 'tuple'>.
"""
def compute_jacobian(*args, **kwargs):
interface = get_interface(*args)
if interface == "autograd":
return _autograd_jacobian(f, argnum=argnums)(*args, **kwargs)
if interface == "jax":
import jax
return jax.jacobian(f, argnums=argnums)(*args, **kwargs)
if interface == "torch":
return _torch_jac(f, argnums, args, kwargs)
if interface == "tensorflow":
return _tensorflow_jac(f, argnums, args, kwargs)
raise ValueError(f"Interface {interface} is not differentiable.")
return compute_jacobian
_modules/pennylane/math/grad
Download Python script
Download Notebook
View on GitHub