diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 097ac57e4d1..830471f991f 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -327,10 +327,13 @@ [(#2063)](https://github.com/PennyLaneAI/pennylane/pull/2063) * Added a new `multi_dispatch` decorator that helps ease the definition of new functions - inside PennyLane. We can decorate the function, indicating the arguments that are - tensors handled by the interface: + inside PennyLane. The decorator is used throughout the math module, demonstrating use cases. [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084) - + [(#2096)](https://github.com/PennyLaneAI/pennylane/pull/2096) + + We can decorate a function, indicating the arguments that are + tensors handled by the interface: + ```pycon >>> @qml.math.multi_dispatch(argnum=[0, 1]) ... def some_function(tensor1, tensor2, option, like): diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py index 5616bdbe863..6b4708e136c 100644 --- a/pennylane/devices/default_qubit.py +++ b/pennylane/devices/default_qubit.py @@ -599,7 +599,7 @@ def density_matrix(self, wires): # Return the full density matrix by using numpy tensor product if wires == self.wires: - density_matrix = self._tensordot(state, self._conj(state), 0) + density_matrix = self._tensordot(state, self._conj(state), axes=0) density_matrix = self._reshape(density_matrix, (2 ** len(wires), 2 ** len(wires))) return density_matrix diff --git a/pennylane/math/is_independent.py b/pennylane/math/is_independent.py index 5a2bea1e67e..302b5bb7792 100644 --- a/pennylane/math/is_independent.py +++ b/pennylane/math/is_independent.py @@ -203,7 +203,7 @@ def _get_random_args(args, interface, num, seed, bounds): if interface == "autograd": # Mark the arguments as trainable with Autograd - rnd_args = pnp.array(rnd_args, requires_grad=True) + rnd_args = [tuple(pnp.array(a, requires_grad=True) for a in arg) for arg in rnd_args] return rnd_args diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index a61288dbd2a..4b8a2d220cb 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -102,8 +102,9 @@ def multi_dispatch(argnum=None, tensor_list=None): to dispatch (i.e., the arguments that are tensors handled by an interface). If ``None``, dispatch over all arguments. tensor_lists (list[int]): a list of integers indicating which indices - in ``argnum`` are expected to be lists of tensors. - If ``None``, this option is ignored. + in ``argnum`` are expected to be lists of tensors. If an argument + marked as tensor list is not a ``tuple`` or ``list``, it is treated + as if it was not marked as tensor list. If ``None``, this option is ignored. Returns: func: A wrapped version of the function, which will automatically attempt @@ -163,7 +164,9 @@ def wrapper(*args, **kwargs): dispatch_args = [] for a in argnums: - if a in tensor_lists: + # Only use extend if the marked argument really + # is a (native) python Sequence + if a in tensor_lists and isinstance(args[a], (list, tuple)): dispatch_args.extend(args[a]) else: dispatch_args.append(args[a]) @@ -179,7 +182,8 @@ def wrapper(*args, **kwargs): return decorator -def block_diag(values): +@multi_dispatch(argnum=[0], tensor_list=[0]) +def block_diag(values, like=None): """Combine a sequence of 2D tensors to form a block diagonal tensor. Args: @@ -203,12 +207,12 @@ def block_diag(values): [ 0, 0, -1, -6, -3, 0], [ 0, 0, 0, 0, 0, 5]]) """ - interface = _multi_dispatch(values) - values = np.coerce(values, like=interface) - return np.block_diag(values, like=interface) + values = np.coerce(values, like=like) + return np.block_diag(values, like=like) -def concatenate(values, axis=0): +@multi_dispatch(argnum=[0], tensor_list=[0]) +def concatenate(values, axis=0, like=None): """Concatenate a sequence of tensors along the specified axis. .. warning:: @@ -235,9 +239,7 @@ def concatenate(values, axis=0): """ - interface = _multi_dispatch(values) - - if interface == "torch": + if like == "torch": import torch if axis is None: @@ -248,16 +250,17 @@ def concatenate(values, axis=0): else: values = [torch.as_tensor(t) for t in values] - if interface == "tensorflow" and axis is None: + if like == "tensorflow" and axis is None: # flatten and then concatenate zero'th dimension # to reproduce numpy's behaviour values = [np.flatten(np.array(t)) for t in values] axis = 0 - return np.concatenate(values, axis=axis, like=interface) + return np.concatenate(values, axis=axis, like=like) -def diag(values, k=0): +@multi_dispatch(argnum=[0], tensor_list=[0]) +def diag(values, k=0, like=None): """Construct a diagonal tensor from a list of scalars. Args: @@ -291,15 +294,14 @@ def diag(values, k=0): [0.0000, 0.0000, 0.2000], [0.0000, 0.0000, 0.0000]]) """ - interface = _multi_dispatch(values) - if isinstance(values, (list, tuple)): - values = np.stack(np.coerce(values, like=interface), like=interface) + values = np.stack(np.coerce(values, like=like), like=like) - return np.diag(values, k=k, like=interface) + return np.diag(values, k=k, like=like) -def dot(tensor1, tensor2): +@multi_dispatch(argnum=[0, 1]) +def dot(tensor1, tensor2, like=None): """Returns the matrix or dot product of two tensors. * If both tensors are 0-dimensional, elementwise multiplication @@ -323,34 +325,34 @@ def dot(tensor1, tensor2): Returns: tensor_like: the matrix or dot product of two tensors """ - interface = _multi_dispatch([tensor1, tensor2]) - x, y = np.coerce([tensor1, tensor2], like=interface) + x, y = np.coerce([tensor1, tensor2], like=like) - if interface == "torch": + if like == "torch": if x.ndim == 0 and y.ndim == 0: return x * y if x.ndim <= 2 and y.ndim <= 2: return x @ y - return np.tensordot(x, y, axes=[[-1], [-2]], like=interface) + return np.tensordot(x, y, axes=[[-1], [-2]], like=like) - if interface == "tensorflow": + if like == "tensorflow": if len(np.shape(x)) == 0 and len(np.shape(y)) == 0: return x * y if len(np.shape(y)) == 1: - return np.tensordot(x, y, axes=[[-1], [0]], like=interface) + return np.tensordot(x, y, axes=[[-1], [0]], like=like) if len(np.shape(x)) == 2 and len(np.shape(y)) == 2: return x @ y - return np.tensordot(x, y, axes=[[-1], [-2]], like=interface) + return np.tensordot(x, y, axes=[[-1], [-2]], like=like) - return np.dot(x, y, like=interface) + return np.dot(x, y, like=like) -def tensordot(tensor1, tensor2, axes=None): +@multi_dispatch(argnum=[0, 1]) +def tensordot(tensor1, tensor2, axes=None, like=None): """Returns the tensor product of two tensors. In general ``axes`` specifies either the set of axes for both tensors that are contracted (with the first/second entry of ``axes`` @@ -376,11 +378,12 @@ def tensordot(tensor1, tensor2, axes=None): Returns: tensor_like: the tensor product of the two input tensors """ - interface = _multi_dispatch([tensor1, tensor2]) - return np.tensordot(tensor1, tensor2, axes=axes, like=interface) + tensor1, tensor2 = np.coerce([tensor1, tensor2], like=like) + return np.tensordot(tensor1, tensor2, axes=axes, like=like) -def get_trainable_indices(values): +@multi_dispatch(argnum=[0], tensor_list=[0]) +def get_trainable_indices(values, like=None): """Returns a set containing the trainable indices of a sequence of values. @@ -403,10 +406,9 @@ def get_trainable_indices(values): tensor(0.0899685, requires_grad=True) """ trainable = requires_grad - interface = _multi_dispatch(values) trainable_params = set() - if interface == "jax": + if like == "jax": import jax if not any(isinstance(v, jax.core.Tracer) for v in values): @@ -420,7 +422,7 @@ def get_trainable_indices(values): trainable = requires_grad for idx, p in enumerate(values): - if trainable(p, interface=interface): + if trainable(p, interface=like): trainable_params.add(idx) return trainable_params @@ -474,8 +476,7 @@ def safe_squeeze(tensor, axis=None, exclude_axis=None): or not excluded and that have size 1. If no axes are specified or excluded, all axes are attempted to be squeezed. """ - interface = _multi_dispatch([tensor]) - if interface == "tensorflow": + if get_interface(tensor) == "tensorflow": from tensorflow.python.framework.errors_impl import InvalidArgumentError exception = InvalidArgumentError @@ -508,7 +509,8 @@ def safe_squeeze(tensor, axis=None, exclude_axis=None): return tensor -def stack(values, axis=0): +@multi_dispatch(argnum=[0], tensor_list=[0]) +def stack(values, axis=0, like=None): """Stack a sequence of tensors along the specified axis. .. warning:: @@ -537,9 +539,8 @@ def stack(values, axis=0): [1.00e-01, 2.00e-01, 3.00e-01], [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)> """ - interface = _multi_dispatch(values) - values = np.coerce(values, like=interface) - return np.stack(values, axis=axis, like=interface) + values = np.coerce(values, like=like) + return np.stack(values, axis=axis, like=like) def where(condition, x=None, y=None): @@ -612,7 +613,8 @@ def where(condition, x=None, y=None): return np.where(condition, x, y, like=_multi_dispatch([condition, x, y])) -def frobenius_inner_product(A, B, normalize=False): +@multi_dispatch(argnum=[0, 1]) +def frobenius_inner_product(A, B, normalize=False, like=None): r"""Frobenius inner product between two matrices. .. math:: @@ -637,8 +639,7 @@ def frobenius_inner_product(A, B, normalize=False): >>> qml.math.frobenius_inner_product(A, B) 3.091948202943376 """ - interface = _multi_dispatch([A, B]) - A, B = np.coerce([A, B], like=interface) + A, B = np.coerce([A, B], like=like) inner_product = np.sum(A * B) @@ -649,6 +650,7 @@ def frobenius_inner_product(A, B, normalize=False): return inner_product +@multi_dispatch(argnum=[0, 2]) def scatter_element_add(tensor, index, value, like=None): """In-place addition of a multidimensional value over various indices of a tensor. @@ -682,8 +684,7 @@ def scatter_element_add(tensor, index, value, like=None): if len(np.shape(tensor)) == 0 and index == (): return tensor + value - interface = like or _multi_dispatch([tensor, value]) - return np.scatter_element_add(tensor, index, value, like=interface) + return np.scatter_element_add(tensor, index, value, like=like) def unwrap(values, max_depth=None): diff --git a/tests/math/test_basic_math.py b/tests/math/test_basic_math.py index 9a098eb9fea..a592749a283 100644 --- a/tests/math/test_basic_math.py +++ b/tests/math/test_basic_math.py @@ -115,7 +115,7 @@ def test_frobenius_inner_product(self, A, B, normalize, expected): def test_frobenius_inner_product_gradient(self): """Test that the calculated gradient is correct.""" - A = np.array([[1.0, 2.3], [-1.3, 2.4]]) + A = onp.array([[1.0, 2.3], [-1.3, 2.4]]) B = torch.autograd.Variable(torch.randn(2, 2).type(torch.float), requires_grad=True) result = fn.frobenius_inner_product(A, B) result.backward() diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py index 2be238b7f2b..8434d7b7313 100644 --- a/tests/math/test_functions.py +++ b/tests/math/test_functions.py @@ -13,6 +13,7 @@ # limitations under the License. """Unit tests for the TensorBox functional API in pennylane.fn.fn """ +from functools import partial import itertools import numpy as onp import pytest @@ -322,7 +323,8 @@ class TestConcatenate: """Tests for the concatenate function""" def test_concatenate_array(self): - """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension""" + """Test that concatenate, called without the axis arguments, + concatenates across the 0th dimension""" t1 = [0.6, 0.1, 0.6] t2 = np.array([0.1, 0.2, 0.3]) t3 = onp.array([5.0, 8.0, 101.0]) @@ -332,7 +334,8 @@ def test_concatenate_array(self): assert np.all(res == np.concatenate([t1, t2, t3])) def test_concatenate_jax(self): - """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension""" + """Test that concatenate, called without the axis arguments, + concatenates across the 0th dimension""" t1 = jnp.array([5.0, 8.0, 101.0]) t2 = jnp.array([0.6, 0.1, 0.6]) t3 = jnp.array([0.1, 0.2, 0.3]) @@ -340,8 +343,9 @@ def test_concatenate_jax(self): res = fn.concatenate([t1, t2, t3]) assert jnp.all(res == jnp.concatenate([t1, t2, t3])) - def test_stack_tensorflow(self): - """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension""" + def test_concatenate_tensorflow(self): + """Test that concatenate, called without the axis arguments, + concatenates across the 0th dimension""" t1 = tf.constant([0.6, 0.1, 0.6]) t2 = tf.Variable([0.1, 0.2, 0.3]) t3 = onp.array([5.0, 8.0, 101.0]) @@ -350,8 +354,9 @@ def test_stack_tensorflow(self): assert isinstance(res, tf.Tensor) assert np.all(res.numpy() == np.concatenate([t1.numpy(), t2.numpy(), t3])) - def test_stack_torch(self): - """Test that concatenate, called without the axis arguments, concatenates across the 0th dimension""" + def test_concatenate_torch(self): + """Test that concatenate, called without the axis arguments, + concatenates across the 0th dimension""" t1 = onp.array([5.0, 8.0, 101.0], dtype=np.float64) t2 = torch.tensor([0.6, 0.1, 0.6], dtype=torch.float64) t3 = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float64) @@ -363,7 +368,7 @@ def test_stack_torch(self): @pytest.mark.parametrize( "t1", [onp.array([[1], [2]]), torch.tensor([[1], [2]]), tf.constant([[1], [2]])] ) - def test_stack_axis(self, t1): + def test_concatenate_axis(self, t1): """Test that passing the axis argument allows for concatenating along a different axis""" t2 = onp.array([[3], [4]]) @@ -546,7 +551,9 @@ def test_multidimensional_product(self, t1, t2): class TestTensordotTorch: - """Tests for the tensor product function in torch.""" + """Tests for the tensor product function in torch. + This test is required because the functionality of tensordot for Torch + is being patched in PennyLane, as compared to autoray.""" v1 = torch.tensor([0.1, 0.5, -0.9, 1.0, -4.2, 0.1], dtype=torch.float64) v2 = torch.tensor([4.3, -1.2, 8.2, 0.6, -4.2, -11.0], dtype=torch.float64) @@ -762,6 +769,89 @@ def test_tensordot_torch_tensor_matrix(self, M, expected, axes1, axes2): assert fn.allclose(fn.tensordot(self.T1, M, axes=[axes1, axes2]), expected) +class TestTensordotDifferentiability: + + v0 = np.array([0.1, 5.3, -0.9, 1.1]) + v1 = np.array([0.5, -1.7, -2.9, 0.0]) + v2 = np.array([-0.4, 9.1, 1.6]) + exp_shapes = ((len(v0), len(v2), len(v0)), (len(v0), len(v2), len(v2))) + exp_jacs = (np.zeros(exp_shapes[0]), np.zeros(exp_shapes[1])) + for i in range(len(v0)): + exp_jacs[0][i, :, i] = v2 + for i in range(len(v2)): + exp_jacs[1][:, i, i] = v0 + + def test_autograd(self): + """Tests differentiability of tensordot with Autograd.""" + v0 = np.array(self.v0, requires_grad=True) + v1 = np.array(self.v1, requires_grad=True) + v2 = np.array(self.v2, requires_grad=True) + + # Test inner product + jac = qml.jacobian(partial(fn.tensordot, axes=[0, 0]), argnum=(0, 1))(v0, v1) + assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0])) + + # Test outer product + jac = qml.jacobian(partial(fn.tensordot, axes=0), argnum=(0, 1))(v0, v2) + assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1]) + assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1]) + + def test_torch(self): + """Tests differentiability of tensordot with Torch.""" + jac_fn = torch.autograd.functional.jacobian + + v0 = torch.tensor(self.v0, requires_grad=True, dtype=torch.float64) + v1 = torch.tensor(self.v1, requires_grad=True, dtype=torch.float64) + v2 = torch.tensor(self.v2, requires_grad=True, dtype=torch.float64) + + # Test inner product + jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), (v0, v1)) + assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0])) + + # Test outer product + jac = jac_fn(partial(fn.tensordot, axes=0), (v0, v2)) + assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1]) + assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1]) + + def test_jax(self): + """Tests differentiability of tensordot with JAX.""" + jac_fn = jax.jacobian + + v0 = jnp.array(self.v0) + v1 = jnp.array(self.v1) + v2 = jnp.array(self.v2) + + # Test inner product + jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), argnums=(0, 1))(v0, v1) + assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0])) + + # Test outer product + jac = jac_fn(partial(fn.tensordot, axes=0), argnums=(0, 1))(v0, v2) + assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1]) + assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1]) + + def test_tensorflow(self): + """Tests differentiability of tensordot with TensorFlow.""" + + def jac_fn(func, args): + with tf.GradientTape() as tape: + out = func(*args) + return tape.jacobian(out, args) + + v0 = tf.Variable(self.v0, dtype=tf.float64) + v1 = tf.Variable(self.v1, dtype=tf.float64) + v2 = tf.Variable(self.v2, dtype=tf.float64) + + # Test inner product + jac = jac_fn(partial(fn.tensordot, axes=[[0], [0]]), (v0, v1)) + assert all(fn.allclose(jac[i], _v) for i, _v in enumerate([v1, v0])) + + # Test outer product + jac = jac_fn(partial(fn.tensordot, axes=0), (v0, v2)) + assert all(fn.shape(jac[i]) == self.exp_shapes[i] for i in [0, 1]) + assert all(fn.allclose(jac[i], self.exp_jacs[i]) for i in [0, 1]) + + # the following test data is of the form # [original shape, axis to expand, new shape] expand_dims_test_data = [ @@ -1362,14 +1452,14 @@ def test_array(self): x = np.array(self.x, requires_grad=True) y = np.array(self.y, requires_grad=True) - def cost(weights): + def cost(*weights): return fn.scatter_element_add(weights[0], self.index, weights[1] ** 2) - res = cost([x, y]) + res = cost(x, y) assert isinstance(res, np.ndarray) assert fn.allclose(res, self.expected_val) - grad = qml.grad(lambda weights: cost(weights)[self.index[0], self.index[1]])([x, y]) + grad = qml.grad(lambda *weights: cost(*weights)[self.index[0], self.index[1]])(x, y) assert fn.allclose(grad[0], self.expected_grad_x) assert fn.allclose(grad[1], self.expected_grad_y) @@ -1492,20 +1582,20 @@ def test_array(self): x = np.array(self.x, requires_grad=True) y = np.array(self.y, requires_grad=True) - def cost(weights): + def cost(*weights): return fn.scatter_element_add( weights[0], self.indices, [fn.sin(weights[1] / 2), weights[1] ** 2] ) - res = cost([x, y]) + res = cost(x, y) assert isinstance(res, np.ndarray) assert fn.allclose(res, self.expected_val) scalar_cost = ( - lambda weights: cost(weights)[self.indices[0][0], self.indices[1][0]] - + cost(weights)[self.indices[0][1], self.indices[1][1]] + lambda *weights: cost(*weights)[self.indices[0][0], self.indices[1][0]] + + cost(*weights)[self.indices[0][1], self.indices[1][1]] ) - grad = qml.grad(scalar_cost)([x, y]) + grad = qml.grad(scalar_cost)(x, y) assert fn.allclose(grad[0], self.expected_grad_x) assert fn.allclose(grad[1], self.expected_grad_y) @@ -1871,7 +1961,7 @@ def test_autograd(self): np.array([[x, 1.2 * y], [x ** 2 - y / 3, -x / y]]), ] f = lambda x, y: fn.block_diag(tensors(x, y)) - x, y = 0.2, 1.5 + x, y = np.array([0.2, 1.5], requires_grad=True) res = qml.jacobian(f)(x, y) exp = self.expected(x, y) assert fn.allclose(res[0], exp[0]) @@ -2029,12 +2119,16 @@ def cost_fn(*params): unwrapped_params = qml.math.unwrap(params) return np.sum(np.sin(params[0] * params[2])) + params[1] - values = [onp.array([0.1, 0.2]), np.tensor(0.1, dtype=np.float64), np.tensor([0.5, 0.2])] - grad = qml.grad(cost_fn)(*values) + values = [ + onp.array([0.1, 0.2]), + np.tensor(0.1, dtype=np.float64, requires_grad=True), + np.tensor([0.5, 0.2], requires_grad=True), + ] + grad = qml.grad(cost_fn, argnum=[1, 2])(*values) expected = [np.array([0.1, 0.2]), 0.1, np.array([0.5, 0.2])] assert all(np.allclose(a, b) for a, b in zip(unwrapped_params, expected)) - assert all(not isinstance(a, ArrayBox) for a in unwrapped_params) + assert not any(isinstance(a, ArrayBox) for a in unwrapped_params) def test_autograd_unwrapping_backward_nested(self): """Test that a sequence of Autograd values is properly unwrapped diff --git a/tests/math/test_is_independent.py b/tests/math/test_is_independent.py index 6953b584df8..cc2af6e6a6f 100644 --- a/tests/math/test_is_independent.py +++ b/tests/math/test_is_independent.py @@ -56,8 +56,8 @@ lambda x: x if abs(x) < 1e-5 else 0.0, # x*delta for x=0 is okay lambda x: 1.0 if x > 0 else 0.0, # Heaviside is okay numerically lambda x: 1.0 if x > 0 else 0.0, # Heaviside is okay numerically - lambda x: qml.math.log(1 + qml.math.exp(1000.0 * x)) / 1000.0, # Softplus is okay - lambda x: qml.math.log(1 + qml.math.exp(1000.0 * x)) / 1000.0, # Softplus is okay + lambda x: qml.math.log(1 + qml.math.exp(100.0 * x)) / 100.0, # Softplus is okay + lambda x: qml.math.log(1 + qml.math.exp(100.0 * x)) / 100.0, # Softplus is okay ] args_dependent_lambdas = [ @@ -167,7 +167,7 @@ def dependent_circuit(x, y, z): dependent_circuit, np.array, lambda x: np.array(x * 0.0), - lambda x: (1 + qml.math.tanh(1000 * x)) / 2, + lambda x: (1 + qml.math.tanh(100 * x)) / 2, *dependent_lambdas, ] @@ -519,7 +519,7 @@ def dependent_circuit(x, y, z): dependent_functions = [ dependent_circuit, - torch.tensor, + torch.as_tensor, lambda x: (1 + qml.math.tanh(1000 * x)) / 2, *dependent_lambdas, ] diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_dispatch.py similarity index 100% rename from tests/math/test_multi_disptach.py rename to tests/math/test_multi_dispatch.py