# Copyright 2019 The TensorNetwork Authors
#
# 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.
# pylint: disable=line-too-long
from typing import Optional, Any, Sequence, Tuple, Type, Callable, List
from typing import Union
from tensornetwork.backends import abstract_backend
from tensornetwork.backends.tensorflow import decompositions
import functools as fct
import operator as op
# This might seem bad, but pytype treats tf.Tensor as Any anyway, so
# we don't actually lose anything by doing this.
import numpy as np
Tensor = Any
# pylint: disable=abstract-method
class TensorFlowBackend(abstract_backend.AbstractBackend):
"""See base_backend.BaseBackend for documentation."""
def __init__(self) -> None:
# pylint: disable=global-variable-undefined
global tf
super().__init__()
try:
# pylint: disable=import-outside-toplevel
import tensorflow
except ImportError as err:
raise ImportError("Tensorflow not installed, please switch to a "
"different backend or install Tensorflow.") from err
tf = tensorflow
self.name = "tensorflow"
def tensordot(self, a: Tensor, b: Tensor,
axes: Union[int, Sequence[Sequence[int]]]) -> Tensor:
return tf.tensordot(a, b, axes)
[docs] def reshape(self, tensor: Tensor, shape: Tensor) -> Tensor:
return tf.reshape(tensor, shape)
[docs] def transpose(self, tensor, perm=None) -> Tensor:
return tf.transpose(tensor, perm)
[docs] def slice(self, tensor: Tensor, start_indices: Tuple[int, ...],
slice_sizes: Tuple[int, ...]) -> Tensor:
if len(start_indices) != len(slice_sizes):
raise ValueError("Lengths of start_indices and slice_sizes must be"
"identical.")
return tf.slice(tensor, start_indices, slice_sizes)
def svd(
self,
tensor: Tensor,
pivot_axis: int = -1,
max_singular_values: Optional[int] = None,
max_truncation_error: Optional[float] = None,
relative: Optional[bool] = False
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
return decompositions.svd(
tf,
tensor,
pivot_axis,
max_singular_values,
max_truncation_error,
relative=relative)
def qr(self, tensor: Tensor, pivot_axis: int = -1,
non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
return decompositions.qr(tf, tensor, pivot_axis, non_negative_diagonal)
def rq(self, tensor: Tensor, pivot_axis: int = -1,
non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
return decompositions.rq(tf, tensor, pivot_axis, non_negative_diagonal)
[docs] def shape_concat(self, values: Tensor, axis: int) -> Tensor:
return tf.concat(values, axis)
[docs] def shape_tensor(self, tensor: Tensor) -> Tensor:
return tf.shape(tensor)
[docs] def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]:
return tuple(tensor.shape.as_list())
[docs] def sparse_shape(self, tensor: Tensor) -> Tuple[Optional[int], ...]:
return self.shape_tuple(tensor)
[docs] def shape_prod(self, values: Tensor) -> Tensor:
return tf.reduce_prod(values)
[docs] def sqrt(self, tensor: Tensor) -> Tensor:
return tf.sqrt(tensor)
[docs] def convert_to_tensor(self, tensor: Tensor) -> Tensor:
result = tf.convert_to_tensor(tensor)
return result
def outer_product(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return tf.tensordot(tensor1, tensor2, 0)
#pylint: disable=unused-argument
[docs] def einsum(self,
expression: str,
*tensors: Tensor,
optimize: bool = True) -> Tensor:
return tf.einsum(expression, *tensors)
[docs] def norm(self, tensor: Tensor) -> Tensor:
return tf.linalg.norm(tensor)
def eye(self,
N: int,
dtype: Optional[Type[np.number]] = None,
M: Optional[int] = None) -> Tensor:
dtype = dtype if dtype is not None else tf.float64
return tf.eye(num_rows=N, num_columns=M, dtype=dtype)
def ones(self,
shape: Tuple[int, ...],
dtype: Optional[Type[np.number]] = None) -> Tensor:
dtype = dtype if dtype is not None else tf.float64
return tf.ones(shape=shape, dtype=dtype)
def zeros(self,
shape: Tuple[int, ...],
dtype: Optional[Type[np.number]] = None) -> Tensor:
dtype = dtype if dtype is not None else tf.float64
return tf.zeros(shape, dtype=dtype)
[docs] def randn(self,
shape: Tuple[int, ...],
dtype: Optional[Type[np.number]] = None,
seed: Optional[int] = None) -> Tensor:
if seed:
tf.random.set_seed(seed)
dtype = dtype if dtype is not None else tf.float64
if (dtype is tf.complex128) or (dtype is tf.complex64):
return tf.complex(
tf.random.normal(shape=shape, dtype=dtype.real_dtype),
tf.random.normal(shape=shape, dtype=dtype.real_dtype))
return tf.random.normal(shape=shape, dtype=dtype)
[docs] def conj(self, tensor: Tensor) -> Tensor:
return tf.math.conj(tensor)
[docs] def eigh(self, matrix: Tensor) -> Tuple[Tensor, Tensor]:
return tf.linalg.eigh(matrix)
[docs] def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return tensor1 + tensor2
[docs] def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return tensor1 - tensor2
[docs] def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return tensor1 * tensor2
[docs] def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
return tensor1 / tensor2
[docs] def index_update(self, tensor: Tensor, mask: Tensor,
assignee: Tensor) -> Tensor:
#returns a copy (unfortunately)
return tf.where(mask, assignee, tensor)
[docs] def inv(self, matrix: Tensor) -> Tensor:
if len(matrix.shape) > 2:
raise ValueError("input to tensorflow backend method `inv` has shape {}. "
"Only matrices are supported.".format(tf.shape(matrix)))
return tf.linalg.inv(matrix)
[docs] def broadcast_right_multiplication(self, tensor1: Tensor,
tensor2: Tensor) -> Tensor:
if len(tensor2.shape) != 1:
raise ValueError("only order-1 tensors are allowed for `tensor2`, "
"found `tensor2.shape = {}`".format(tf.shape(tensor2)))
return tensor1 * tensor2
[docs] def broadcast_left_multiplication(self, tensor1: Tensor,
tensor2: Tensor) -> Tensor:
if len(tensor1.shape) != 1:
raise ValueError("only order-1 tensors are allowed for `tensor1`,"
" found `tensor1.shape = {}`".format(tf.shape(tensor1)))
t1_broadcast_shape = self.shape_concat(
[self.shape_tensor(tensor1), [1] * (len(tensor2.shape) - 1)], axis=-1)
return tensor2 * self.reshape(tensor1, t1_broadcast_shape)
def sin(self, tensor: Tensor) -> Tensor:
return tf.math.sin(tensor)
def cos(self, tensor: Tensor) -> Tensor:
return tf.math.cos(tensor)
[docs] def exp(self, tensor: Tensor) -> Tensor:
return tf.math.exp(tensor)
[docs] def log(self, tensor: Tensor) -> Tensor:
return tf.math.log(tensor)
def expm(self, matrix: Tensor) -> Tensor:
if len(matrix.shape) != 2:
raise ValueError("input to tensorflow backend method `expm` has shape {}."
" Only matrices are supported.".format(matrix.shape))
if matrix.shape[0] != matrix.shape[1]:
raise ValueError("input to tensorflow backend method `expm` only supports"
"N*N matrix, {x}*{y} matrix is given".format(
x=matrix.shape[0], y=matrix.shape[1]))
return tf.linalg.expm(matrix)
def jit(self, fun: Callable, *args: List, **kwargs: dict) -> Callable:
# tf.function is slow and bad.
return fun
[docs] def sum(self,
tensor: Tensor,
axis: Optional[Sequence[int]] = None,
keepdims: bool = False) -> Tensor:
return tf.math.reduce_sum(tensor, axis=axis, keepdims=keepdims)
def matmul(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
if (tensor1.ndim <= 1) or (tensor2.ndim <= 1):
raise ValueError("inputs to `matmul` have to be a tensors of order > 1,")
return tf.matmul(tensor1, tensor2)
[docs] def diagonal(self, tensor: Tensor, offset: int = 0, axis1: int = -2,
axis2: int = -1) -> Tensor:
"""Return specified diagonals.
If tensor is 2-D, returns the diagonal of tensor with the given offset,
i.e., the collection of elements of the form a[i, i+offset].
If a has more than two dimensions, then the axes specified by
axis1 and axis2 are used to determine the 2-D sub-array whose diagonal is
returned. The shape of the resulting array can be determined by removing
axis1 and axis2 and appending an index to the right equal to the size of the
resulting diagonals.
This function only extracts diagonals. If you
wish to create diagonal matrices from vectors, use diagflat.
Args:
tensor: A tensor.
offset: Offset of the diagonal from the main diagonal.
axis1, axis2: Axis to be used as the first/second axis of the 2D
sub-arrays from which the diagonals should be taken.
Defaults to second-last and last axis (note this
differs from the NumPy defaults).
These arguments are not supported in the TensorFlow
backend and an error will be raised if they are
specified.
Returns:
array_of_diagonals: A dim = min(1, tensor.ndim - 2) tensor storing
the batched diagonals.
"""
if axis1 != -2 or axis2 != -1:
errstr = (f"axis1={axis1}, axis2={axis2} must be -2, -1 (the defaults)"
f"with TensorFlow backend.")
raise NotImplementedError(errstr)
#pylint: disable=unexpected-keyword-arg
return tf.linalg.diag_part(tensor, k=offset)
[docs] def diagflat(self, tensor: Tensor, k: int = 0) -> Tensor:
""" Flattens tensor and creates a new matrix of zeros with its elements
on the k'th diagonal.
Args:
tensor: A tensor.
k : The diagonal upon which to place its elements.
Returns:
tensor: A new tensor with all zeros save the specified diagonal.
"""
#pylint: disable=unexpected-keyword-arg
return tf.linalg.diag(tensor, k=k)
[docs] def trace(self, tensor: Tensor, offset: int = 0, axis1: int = -2,
axis2: int = -1) -> Tensor:
"""Return summed entries along diagonals.
If tensor is 2-D, the sum is over the
diagonal of tensor with the given offset,
i.e., the collection of elements of the form a[i, i+offset].
If a has more than two dimensions, then the axes specified by
axis1 and axis2 are used to determine the 2-D sub-array whose diagonal is
summed.
Args:
tensor: A tensor.
offset: Offset of the diagonal from the main diagonal.
This argument is not supported in the TensorFlow
backend and an error will be raised if they are
specified.
axis1, axis2: Axis to be used as the first/second axis of the 2D
sub-arrays from which the diagonals should be taken.
Defaults to first/second axis.
These arguments are not supported in the TensorFlow
backend and an error will be raised if they are
specified.
Returns:
array_of_diagonals: The batched summed diagonals.
"""
if offset != 0:
errstr = (f"offset = {offset} must be 0 (the default)"
f"with TensorFlow backend.")
raise NotImplementedError(errstr)
if axis1 == axis2:
raise ValueError(f"axis1 = {axis1} cannot equal axis2 = {axis2}")
N = len(tensor.shape)
if N > 25:
raise ValueError(f"Currently only tensors with ndim <= 25 can be traced"
f"in the TensorFlow backend (yours was {N})")
if axis1 < 0:
axis1 = N+axis1
if axis2 < 0:
axis2 = N+axis2
inds = list(map(chr, range(98, 98+N)))
indsout = [i for n, i in enumerate(inds) if n not in (axis1, axis2)]
inds[axis1] = 'a'
inds[axis2] = 'a'
return tf.einsum(''.join(inds) + '->' +''.join(indsout), tensor)
def abs(self, tensor: Tensor) -> Tensor:
"""
Returns the elementwise absolute value of tensor.
Args:
tensor: An input tensor.
Returns:
tensor: Its elementwise absolute value.
"""
return tf.math.abs(tensor)
[docs] def sign(self, tensor: Tensor) -> Tensor:
"""
Returns an elementwise tensor with entries
y[i] = 1, 0, -1 where tensor[i] > 0, == 0, and < 0 respectively.
For complex input the behaviour of this function may depend on the backend.
The TensorFlow version returns y[i] = x[i] / abs(x[i]).
Args:
tensor: The input tensor.
"""
return tf.math.sign(tensor)
[docs] def item(self, tensor):
numel = 0
if len(tensor.shape) > 0:
numel = fct.reduce(op.mul, tensor.shape)
if numel != 1:
raise ValueError(f"expected tensor with one element, "
f"got {tensor.shape}")
if numel == 1:
return tensor[0]
return tensor
[docs] def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
"""
Returns the exponentiation of tensor a raised to b.
If b is a tensor, then the exponentiation is element-wise
between the two tensors, with a as the base and b as the power.
Note that a and b must be broadcastable to the same shape if
b is a tensor.
If b is a scalar, then the exponentiation is each value in a
raised to the power of b.
Args:
a: The tensor containing the bases.
b: The tensor containing the powers; or a single scalar as the power.
Returns:
The tensor that is each element of a raised to the
power of b. Note that the shape of the returned tensor
is that produced by the broadcast of a and b.
"""
return tf.math.pow(a, b)
[docs] def eps(self, dtype: Type[np.number]) -> float:
"""
Return machine epsilon for given `dtype`
Args:
dtype: A dtype.
Returns:
float: Machine epsilon.
"""
return tf.experimental.numpy.finfo(dtype).eps