# Source code for pennylane.numpy.wrapper

# 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

# 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 provides the PennyLane wrapper functions for modifying NumPy,
such that it accepts the PennyLane :class:~.tensor class.
"""
from collections.abc import Sequence
import functools

from autograd import numpy as _np

from .tensor import tensor

[docs]def extract_tensors(x):
"""Iterate through an iterable, and extract any PennyLane
tensors that appear.

Args:
x (.tensor or Sequence): an input tensor or sequence

Yields:
tensor: the next tensor in the sequence. If the input was a single
tensor, than the tensor is yielded and the iterator completes.

**Example**

>>> from pennylane import numpy as np
>>> import numpy as onp
>>> iterator = np.extract_tensors([0.1, np.array(0.1), "string", onp.array(0.5)])
>>> list(iterator)
"""
if isinstance(x, tensor):
# If the item is a tensor, return it
yield x
elif isinstance(x, Sequence) and not isinstance(x, (str, bytes)):
# If the item is a sequence, recursively look through its
# elements for tensors.
# NOTE: we choose to branch on Sequence here and not Iterable,
# as NumPy arrays are not Sequences.
for item in x:
yield from extract_tensors(item)

[docs]def tensor_wrapper(obj):
"""Decorator that wraps callable objects and classes so that they both accept
a requires_grad keyword argument, as well as returning a PennyLane
:class:~.tensor.

Only if the decorated object returns an ndarray is the
output converted to a :class:~.tensor; this avoids superfluous conversion
of scalars and other native-Python types.

.. note::

This wrapper does *not* enable autodifferentiation of the wrapped function,
it merely adds support for :class:~pennylane.numpy.tensor output.

Args:
obj: a callable object or class

**Example**

By default, the ones function provided by Autograd
constructs standard ndarray objects, and does not
permit a requires_grad argument:

>>> from autograd.numpy import ones
>>> ones([2, 2])
array([[1., 1.],
[1., 1.]])
>>> ones([2, 2], requires_grad=True)
TypeError: ones() got an unexpected keyword argument 'requires_grad'

tensor_wrapper both enables construction of :class:~pennylane.numpy.tensor
objects, while also converting the output.

>>> from pennylane import numpy as np
>>> ones = np.tensor_wrapper(ones)
>>> ones([2, 2], requires_grad=True)
tensor([[1., 1.],
"""

@functools.wraps(obj)
def _wrapped(*args, **kwargs):
"""Wrapped NumPy function"""

tensor_kwargs = {}

if "requires_grad" in kwargs:
else:
tensor_args = list(extract_tensors(args))

if tensor_args:
# Unless the user specifies otherwise, if all tensors in the argument
# list are non-trainable, the output is also non-trainable.
# Equivalently: if any tensor is trainable, the output is also trainable.
# NOTE: Use of Python's any results in an infinite recursion,
# and I'm not sure why. Using np.any works fine.
tensor_kwargs["requires_grad"] = _np.any([i.requires_grad for i in tensor_args])

# evaluate the original object
res = obj(*args, **kwargs)

if isinstance(res, _np.ndarray):
# only if the output of the object is a ndarray,
# then convert to a PennyLane tensor
res = tensor(res, **tensor_kwargs)

return res

return _wrapped

[docs]def wrap_arrays(old, new):
"""Loop through an object's symbol table,
wrapping each function with :func:~pennylane.numpy.tensor_wrapper.

This is useful if you would like to wrap **every** function
provided by an imported module.

Args:
old (dict): The symbol table to be wrapped. Note that
callable classes are ignored; only functions are wrapped.
new (dict): The symbol table that contains the wrapped values.

.. seealso:: :func:~pennylane.numpy.tensor_wrapper

**Example**

This function is used to wrap the imported autograd.numpy
module, to enable all functions to support requires_grad
arguments, and to output :class:~pennylane.numpy.tensor objects:

>>> from autograd import numpy as _np
>>> wrap_arrays(_np.__dict__, globals())
"""
for name, obj in old.items():
if callable(obj) and not isinstance(obj, type):
new[name] = tensor_wrapper(obj)


Using PennyLane

Development

API