def test_node_invalid_network(backend): a = np.ones((2, 2)) b = np.ones((2, 3, 4, 2)) c = np.ones((2, 4, 4, 3)) run_tests( Tensor(a, backend=backend), Tensor(b, backend=backend), Tensor(c, backend=backend), backend)
def test_kron(backend, dtype): """ Checks that Tensor.kron() works. """ if (backend == "pytorch" and dtype in (np.complex64, np.complex128)): pytest.skip("pytorch support for complex dtypes is currently poor.") np.random.seed(10) t1 = Tensor(np.random.rand(2, 2).astype(dtype), backend=backend) t2 = Tensor(np.random.rand(3, 3).astype(dtype), backend=backend) res_kron = kron(t1, t2) res_ncon = ncon([t1.array, t2.array], [[-1, -3], [-2, -4]], backend=backend) np.testing.assert_allclose(res_kron.array, res_ncon) mat1 = res_kron.reshape((6, 6)) mat2 = np.kron(t1.array, t2.array) np.testing.assert_allclose(mat1.array, mat2) t1 = Tensor(np.random.rand(2, 2, 2, 2).astype(dtype), backend=backend) t2 = Tensor(np.random.rand(3, 3, 3, 3).astype(dtype), backend=backend) res_kron = kron(t1, t2) res_ncon = ncon([t1.array, t2.array], [[-1, -2, -5, -6], [-3, -4, -7, -8]], backend=backend) np.testing.assert_allclose(res_kron.array, res_ncon) mat1 = res_kron.reshape((36, 36)) mat2 = np.kron( np.array(t1.array).reshape(4, 4), np.array(t2.array).reshape(9, 9)) np.testing.assert_allclose(mat1.array, mat2)
def test_return_type(backend): t1, t2 = np.ones((2, 2)), np.ones((2, 2)) n1, n2 = Tensor(t1, backend=backend), Tensor(t2, backend=backend) result_1 = ncon_interface.ncon([t1, t2], [(-1, 1), (1, -2)], backend=backend) result_2 = ncon_interface.ncon([n1, n2], [(-1, 1), (1, -2)], backend=backend) result_3 = ncon_interface.ncon([n1, t2], [(-1, 1), (1, -2)], backend=backend) assert isinstance(result_1, type(n1.backend.convert_to_tensor(t1))) assert isinstance(result_2, Tensor) assert isinstance(result_3, type(n1.backend.convert_to_tensor(t1)))
def test_node_small_matmul(backend): np.random.seed(10) t1 = np.random.randn(2, 2) t2 = np.random.randn(2, 2) a = Tensor(t1, backend=backend) b = Tensor(t2, backend=backend) res = ncon_interface.ncon([a, b], [(1, -1), (1, -2)], backend=backend) np.testing.assert_allclose(res.array, t1.transpose() @ t2)
def test_node_outer_product_1(backend): t1 = np.array([1, 2, 3]) t2 = np.array([1, 2]) a = Tensor(t1, backend=backend) b = Tensor(t2, backend=backend) res = ncon_interface.ncon([a, b], [(-1, ), (-2, )], backend=backend) np.testing.assert_allclose(res.array, np.kron(t1, t2).reshape((3, 2))) res = ncon_interface.ncon([a, a, a, a], [(1, ), (1, ), (2, ), (2, )], backend=backend) np.testing.assert_allclose(res.array, 196)
def test_node_outer_product_2_mixed_labels(backend): np.random.seed(10) t1 = np.random.rand(10, 100) t2 = np.random.rand(8) a = Tensor(t1, backend=backend) b = Tensor(t2, backend=backend) res = ncon_interface.ncon([a, b], [(-1, '-hi'), ('-ho', )], out_order=['-hi', -1, '-ho'], backend=backend) exp = np.einsum('ij,k->jik', t1, t2) np.testing.assert_allclose(res.array, exp)
def test_kron_raises(backend): np.random.seed(10) t1 = Tensor(np.random.rand(2, 2, 2), backend=backend) t2 = Tensor(np.random.rand(3, 3), backend=backend) with pytest.raises(ValueError, match="tensorA.ndim"): kron(t1, t2) with pytest.raises(ValueError, match="tensorB.ndim"): kron(t2, t1) t1 = Tensor(np.random.rand(2, 2, 2), backend='numpy') t2 = Tensor(np.random.rand(3, 3), backend='tensorflow') with pytest.raises(ValueError, match="kron"): kron(t1, t2)
def test_return_type(backend): t1, t2 = np.ones((2, 2)), np.ones((2, 2)) n1, n2 = Tensor(t1, backend=backend), Tensor(t2, backend=backend) result_1 = ncon_interface.ncon([t1, t2], [(-1, 1), (1, -2)], backend=backend) result_2 = ncon_interface.ncon([n1, n2], [(-1, 1), (1, -2)], backend=backend) result_3 = ncon_interface.ncon([n1, t2], [(-1, 1), (1, -2)], backend=backend) assert isinstance(result_2, Tensor) if isinstance(backend, JaxBackend) or backend == 'jax': pytest.skip('return-type tests for jax not implemented') assert isinstance(result_1, type(n1.backend.convert_to_tensor(t1))) assert isinstance(result_3, type(n1.backend.convert_to_tensor(t1)))
def test_node_order_spec(backend): np.random.seed(10) a = np.random.rand(2, 2) node = Tensor(a, backend=backend) result = ncon_interface.ncon([node, node], [(-1, 1), (1, -2)], out_order=[-1, -2], backend=backend) np.testing.assert_allclose(result.array, a @ a) result = ncon_interface.ncon([node, node], [(-1, 1), (1, -2)], con_order=[1], backend=backend) np.testing.assert_allclose(result.array, a @ a) result = ncon_interface.ncon([node, node], [(-1, 1), (1, -2)], con_order=[1], out_order=[-1, -2], backend=backend) np.testing.assert_allclose(result.array, a @ a) result = ncon_interface.ncon([node, node], [(-1, 1), (1, -2)], con_order=[1], out_order=[-2, -1], backend=backend) np.testing.assert_allclose(result.array, (a @ a).T)
def initialize_tensor(fname: Text, *fargs: Any, backend: Optional[Union[Text, AbstractBackend]] = None, **fkwargs: Any) -> Tensor: """Return a Tensor wrapping data obtained by an initialization function implemented in a backend. The Tensor will have the same shape as the underlying array that function generates, with all Edges dangling. This function is not intended to be called directly, but doing so should be safe enough. Args: fname: Name of the method of backend to call (a string). *fargs: Positional arguments to the initialization method. backend: The backend or its name. **fkwargs: Keyword arguments to the initialization method. Returns: tensor: A Tensor wrapping data generated by (the_backend).fname(*fargs, **fkwargs), with one dangling edge per axis of data. """ if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) func = getattr(backend_obj, fname) data = func(*fargs, **fkwargs) tensor = Tensor(data, backend=backend) return tensor
def outer(tensor1: Tensor, tensor2: Tensor) -> Tensor: """Calculate the outer product of the two given Tensors.""" tensors = [tensor1, tensor2] all_backends_same, errstr = _check_backends(tensors, "outer") if not all_backends_same: raise ValueError(errstr) out_data = tensor1.backend.outer_product(tensor1.array, tensor2.array) return Tensor(out_data, backend=tensor1.backend)
def conj(tensor: Tensor) -> Tensor: """ Return the complex conjugate of `Tensor` Args: Tensor: A Tensor. Returns: The complex conjugated Tensor. """ return tensor.conj()
def transpose(tensor: Tensor, perm: Optional[Sequence[int]] = None) -> Tensor: """ Return a new `Tensor` transposed according to the permutation set by `axes`. By default the axes are reversed. Args: axes: The permutation. If None (default) the index order is reversed. Returns: The transposed `Tensor`. """ return tensor.transpose(perm=perm)
def einsum(expression: Text, *tensors: Tensor, optimize: bool) -> Tensor: """Calculate sum of products of Tensors according to expression.""" all_backends_same, errstr = _check_backends(tensors, "einsum") if not all_backends_same: raise ValueError(errstr) backend = tensors[0].backend arrays = [tensor.array for tensor in tensors] result_data = backend.einsum(expression, *arrays, optimize=optimize) return Tensor(result_data, backend=backend)
def reshape(tensor: Tensor, new_shape: Sequence[int]) -> Tensor: """Reshape Tensor to the given shape. Args: tensor: Tensor to reshape. new_shape: The new shape. Returns: The reshaped Tensor. """ return tensor.reshape(new_shape)
def hconj(tensor: Tensor, perm: Optional[Sequence[int]] = None) -> Tensor: """ The Hermitian conjugated tensor; e.g. the complex conjugate tranposed by the permutation set be `axes`. By default the axes are reversed. Args: tensor: The Tensor to conjugate. axes: The permutation. If None (default) the index order is reversed. Returns: The Hermitian conjugated `Tensor`. """ return tensor.hconj(perm=perm)
def sin(tensor: Tensor) -> Tensor: """ Return sin of `Tensor`. Args: Tensor: A Tensor. Returns: Tensor """ out_array = tensor.backend.sin(tensor.array) return Tensor(out_array, backend=tensor.backend)
def exp(tensor: Tensor) -> Tensor: """ Return elementwise exp of `Tensor`. Args: Tensor: A Tensor. Returns: Tensor """ out_array = tensor.backend.exp(tensor.array) return Tensor(out_array, backend=tensor.backend)
def cos(tensor: Tensor) -> Tensor: """ Return cos of `Tensor`. Args: Tensor: A Tensor. Returns: Tensor """ out_array = tensor.backend.cos(tensor.array) return Tensor(out_array, backend=tensor.backend)
def log(tensor: Tensor) -> Tensor: """ Return elementwise natural logarithm of `Tensor`. Args: Tensor: A Tensor. Returns: Tensor """ out_array = tensor.backend.log(tensor.array) return Tensor(out_array, backend=tensor.backend)
def test_node_contraction(backend): np.random.seed(10) tensor = np.random.randn(2, 2, 2) a = Tensor(tensor, backend=backend) res = ncon_interface.ncon([a, a, a], [(-1, 1, 2), (1, 2, 3), (3, -2, -3)], backend=backend) res_np = tensor.reshape((2, 4)) @ tensor.reshape((4, 2)) @ tensor.reshape( (2, 4)) res_np = res_np.reshape((2, 2, 2)) np.testing.assert_allclose(res.array, res_np)
def test_node_contraction_mixed_labels(backend): np.random.seed(10) tensor = np.random.randn(2, 2, 2) a = Tensor(tensor, backend=backend) res = ncon_interface.ncon([a, a, a], [(-1, 'rick', 2), ('rick', 2, 'morty'), ('morty', -2, -3)], backend=backend) res_np = tensor.reshape((2, 4)) @ tensor.reshape((4, 2)) @ tensor.reshape( (2, 4)) res_np = res_np.reshape((2, 2, 2)) np.testing.assert_allclose(res.array, res_np)
def pivot(tensor: Tensor, pivot_axis: int = -1) -> Tensor: """ Reshapes tensor into a matrix about the pivot_axis. Equivalent to tensor.reshape(prod(tensor.shape[:pivot_axis]), prod(tensor.shape[pivot_axis:])). Args: tensor: The input tensor. pivot_axis: Axis to pivot around. """ backend = tensor.backend result = backend.pivot(tensor.array, pivot_axis=pivot_axis) return Tensor(result, backend=backend)
def expm(matrix: Tensor) -> Tensor: """ Return expm log of `matrix`, matrix exponential. Args: matrix: A tensor. Returns: Tensor """ backend = matrix.backend out = backend.expm(matrix.array) tensor = Tensor(out, backend=backend) return tensor
def inv(matrix: Tensor) -> Tensor: """Compute the matrix inverse of `matrix`. Args: matrix: A matrix. Returns: Tensor: The inverse of `matrix` """ backend = matrix.backend out = backend.inv(matrix.array) tensor = Tensor(out, backend=backend) return tensor
def eigh(matrix: Tensor) -> Tuple[Tensor, Tensor]: """Compute eigenvectors and eigenvalues of a hermitian matrix. Args: matrix: A symetric matrix. Returns: Tensor: The eigenvalues in ascending order. Tensor: The eigenvectors. """ backend = matrix.backend out = backend.eigh(matrix.array) tensors = [Tensor(t, backend=backend) for t in out] return tuple(tensors)
def take_slice(tensor: Tensor, start_indices: Tuple[int, ...], slice_sizes: Tuple[int, ...]) -> Tensor: """Obtains a slice of a Tensor based on start_indices and slice_sizes. Args: Tensor: A Tensor. start_indices: Tuple of integers denoting start indices of slice. slice_sizes: Tuple of integers denoting size of slice along each axis. Returns: The slice, a Tensor. """ sliced = tensor.backend.slice(tensor.array, start_indices, slice_sizes) sliced_tensor = Tensor(sliced, backend=tensor.backend) return sliced_tensor
def diagflat(tensor: Tensor, k: int = 0) -> Tensor: """ Flattens tensor and places its elements at the k'th diagonal of a new (tensor.size + k, tensor.size + k) `Tensor` of zeros. Args: tensor: A Tensor. k : The elements of tensor will be stored at this diagonal. Returns: out : A (tensor.size + k, tensor.size + k) `Tensor` with the elements of tensor on its kth diagonal. """ backend = tensor.backend result = backend.diagflat(tensor.array, k=k) return Tensor(result, backend=backend)
def test_node_order_spec_noninteger(backend): np.random.seed(10) a = np.random.rand(2, 2) exp = a @ a node = Tensor(a, backend=backend) result = ncon_interface.ncon([node, node], [('-o1', 'i'), ('i', '-o2')], con_order=['i'], out_order=['-o1', '-o2'], backend=backend) np.testing.assert_allclose(result.array, exp) result = ncon_interface.ncon([node, node], [('-o1', 'i'), ('i', '-o2')], con_order=['i'], out_order=['-o2', '-o1'], backend=backend) np.testing.assert_allclose(result.array, exp.T)
def trace(tensor: Tensor, offset: int = 0, axis1: int = -2, axis2: int = -1) -> Tensor: """Calculate the sum along diagonal entries of the given Tensor. The entries of the offset`th diagonal of the matrix slice of tensor indexed by (axis1, axis2) are summed. Args: tensor: A Tensor. offset: Offset of the diagonal from the main diagonal. axis1, axis2: Indices of the matrix slice to extract from. Returns: out: The trace. """ backend = tensor.backend result = backend.trace(tensor.array, offset=offset, axis1=axis1, axis2=axis2) return Tensor(result, backend=backend)