tensorcircuit.backends.numpy_backend#

Backend magic inherited from tensornetwork: numpy backend

class tensorcircuit.backends.numpy_backend.NumpyBackend[source]#

Bases: tensornetwork.backends.numpy.numpy_backend.NumPyBackend, tensorcircuit.backends.abstract_backend.ExtendedBackend

see the original backend API at numpy backend

__init__() β†’ None#
abs(a: Any) β†’ Any[source]#

Returns the elementwise absolute value of tensor. :param tensor: An input tensor.

Returns

Its elementwise absolute value.

Return type

tensor

acos(a: Any) β†’ Any[source]#

Return the acos of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

acos of a

Return type

Tensor

acosh(a: Any) β†’ Any[source]#

Return the acosh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

acosh of a

Return type

Tensor

addition(tensor1: Any, tensor2: Any) β†’ Any#

Return the default addition of tensor. A backend can override such implementation. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

Tensor

adjoint(a: Any) β†’ Any#

Return the conjugate and transpose of a tensor a

Parameters

a (Tensor) – Input tensor

Returns

adjoint tensor of a

Return type

Tensor

arange(start: int, stop: Optional[int] = None, step: int = 1) β†’ Any[source]#

Values are generated within the half-open interval [start, stop)

Parameters
  • start (int) – start index

  • stop (Optional[int], optional) – end index, defaults to None

  • step (Optional[int], optional) – steps, defaults to 1

Returns

_description_

Return type

Tensor

argmax(a: Any, axis: int = 0) β†’ Any[source]#

Return the index of maximum of an array an axis.

Parameters
  • a (Tensor) – [description]

  • axis (int) – [description], defaults to 0, different behavior from numpy defaults!

Returns

[description]

Return type

Tensor

argmin(a: Any, axis: int = 0) β†’ Any[source]#

Return the index of minimum of an array an axis.

Parameters
  • a (Tensor) – [description]

  • axis (int) – [description], defaults to 0, different behavior from numpy defaults!

Returns

[description]

Return type

Tensor

asin(a: Any) β†’ Any[source]#

Return the acos of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

asin of a

Return type

Tensor

asinh(a: Any) β†’ Any[source]#

Return the asinh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

asinh of a

Return type

Tensor

atan(a: Any) β†’ Any[source]#

Return the atan of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

atan of a

Return type

Tensor

atan2(y: Any, x: Any) β†’ Any[source]#

Return the atan of a tensor y/x.

Parameters

a (Tensor) – tensor in matrix form

Returns

atan2 of a

Return type

Tensor

atanh(a: Any) β†’ Any[source]#

Return the atanh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

atanh of a

Return type

Tensor

broadcast_left_multiplication(tensor1: Any, tensor2: Any) β†’ Any#

Perform broadcasting for multiplication of tensor1 onto tensor2, i.e. tensor1 * tensor2`, where tensor2 is an arbitrary tensor and tensor1 is a one-dimensional tensor. The broadcasting is applied to the first index of tensor2. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

The result of multiplying tensor1 onto tensor2.

Return type

Tensor

broadcast_right_multiplication(tensor1: Any, tensor2: Any) β†’ Any#

Perform broadcasting for multiplication of tensor2 onto tensor1, i.e. tensor1 * tensor2`, where tensor1 is an arbitrary tensor and tensor2 is a one-dimensional tensor. The broadcasting is applied to the last index of tensor1. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

The result of multiplying tensor1 onto tensor2.

Return type

Tensor

cast(a: Any, dtype: str) β†’ Any[source]#

Cast the tensor dtype of a a.

Parameters
  • a (Tensor) – tensor

  • dtype (str) – β€œfloat32”, β€œfloat64”, β€œcomplex64”, β€œcomplex128”

Returns

a of new dtype

Return type

Tensor

cholesky(tensor: Any, pivot_axis: int = - 1, non_negative_diagonal: bool = False) β†’ Tuple[Any, Any]#
concat(a: Sequence[Any], axis: int = 0) β†’ Any[source]#

Join a sequence of arrays along an existing axis.

Parameters
  • a (Sequence[Tensor]) – [description]

  • axis (int, optional) – [description], defaults to 0

cond(pred: bool, true_fun: Callable[[], Any], false_fun: Callable[[], Any]) β†’ Any[source]#

The native cond for XLA compiling, wrapper for tf.cond and limited functionality of jax.lax.cond.

Parameters
  • pred (bool) – [description]

  • true_fun (Callable[[], Tensor]) – [description]

  • false_fun (Callable[[], Tensor]) – [description]

Returns

[description]

Return type

Tensor

conj(tensor: Any) β†’ Any#

Return the complex conjugate of tensor :param tensor: A tensor.

Returns

Tensor

convert_to_tensor(a: Any) β†’ Any#

Convert a np.array or a tensor to a tensor type for the backend.

coo_sparse_matrix(indices: Any, values: Any, shape: Any) β†’ Any[source]#

Generate the coo format sparse matrix from indices and values, which is the only sparse format supported in different ML backends.

Parameters
  • indices (Tensor) – shape [n, 2] for n non zero values in the returned matrix

  • values (Tensor) – shape [n]

  • shape (Tensor) – Tuple[int, …]

Returns

[description]

Return type

Tensor

coo_sparse_matrix_from_numpy(a: Any) β†’ Any#

Generate the coo format sparse matrix from scipy coo sparse matrix.

Parameters

a (Tensor) – Scipy coo format sparse matrix

Returns

SparseTensor in backend format

Return type

Tensor

copy(a: Any) β†’ Any[source]#

Return the copy of a, matrix exponential.

Parameters

a (Tensor) – tensor in matrix form

Returns

the copy tensor of a

Return type

Tensor

cos(a: Any) β†’ Any[source]#

Return cos of tensor. :param tensor: A tensor.

Returns

Tensor

cosh(a: Any) β†’ Any[source]#

Return the cosh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

cosh of a

Return type

Tensor

cumsum(a: Any, axis: Optional[int] = None) β†’ Any[source]#

Return the cumulative sum of the elements along a given axis.

Parameters
  • a (Tensor) – [description]

  • axis (Optional[int], optional) – The default behavior is the same as numpy, different from tf/torch as cumsum of the flatten 1D array, defaults to None

Returns

[description]

Return type

Tensor

deserialize_tensor(s: str) β†’ Any#

Return a tensor given a serialized tensor string.

Parameters

s – The input string representing a serialized tensor.

Returns

The tensor object represented by the string.

det(a: Any) β†’ Any[source]#

Return the determinant scalar of a tensor a

Parameters

a (Tensor) – Input tensor

Returns

determinant of a

Return type

Tensor

device(a: Any) β†’ str[source]#

get the universal device str for the tensor, in the format of tf

Parameters

a (Tensor) – the tensor

Returns

device str where the tensor lives on

Return type

str

device_move(a: Any, dev: Any) β†’ Any[source]#

move tensor a to device dev

Parameters
  • a (Tensor) – the tensor

  • dev (Any) – device str or device obj in corresponding backend

Returns

the tensor on new device

Return type

Tensor

diagflat(tensor: Any, k: int = 0) β†’ Any#

Flattens tensor and creates a new matrix of zeros with its elements on the k’th diagonal. :param tensor: A tensor. :param k: The diagonal upon which to place its elements.

Returns

A new tensor with all zeros save the specified diagonal.

Return type

tensor

diagonal(tensor: Any, offset: int = 0, axis1: int = - 2, axis2: int = - 1) β†’ Any#

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.

Parameters
  • tensor – A tensor.

  • offset – Offset of the diagonal from the main diagonal.

  • axis1 – 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/last axis.

  • 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/last axis.

Returns

A dim = min(1, tensor.ndim - 2) tensor storing

the batched diagonals.

Return type

array_of_diagonals

divide(tensor1: Any, tensor2: Any) β†’ Any#

Return the default divide of tensor. A backend can override such implementation. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

Tensor

dtype(a: Any) β†’ str[source]#

Obtain dtype string for tensor a

Parameters

a (Tensor) – The tensor

Returns

dtype str, such as β€œcomplex64”

Return type

str

eigh(matrix: Any) β†’ Tuple[Any, Any]#

Compute eigenvectors and eigenvalues of a hermitian matrix.

Parameters

matrix – A symetric matrix.

Returns

The eigenvalues in ascending order. Tensor: The eigenvectors.

Return type

Tensor

eigs(A: Callable, args: Optional[List] = None, initial_state: Optional[Any] = None, shape: Optional[Tuple[int, ...]] = None, dtype: Optional[Type[numpy.number]] = None, num_krylov_vecs: int = 50, numeig: int = 6, tol: float = 1e-08, which: str = 'LR', maxiter: Optional[int] = None) β†’ Tuple[Any, List]#

Arnoldi method for finding the lowest eigenvector-eigenvalue pairs of a linear operator A. If no initial_state is provided then shape and dtype are required so that a suitable initial state can be randomly generated. This is a wrapper for scipy.sparse.linalg.eigs which only supports a subset of the arguments of scipy.sparse.linalg.eigs.

Parameters
  • A – A (sparse) implementation of a linear operator

  • args – A list of arguments to A. A will be called as res = A(initial_state, *args).

  • initial_state – An initial vector for the algorithm. If None, a random initial Tensor is created using the numpy.random.randn method.

  • shape – The shape of the input-dimension of A.

  • dtype – The dtype of the input A. If both no initial_state is provided, a random initial state with shape shape and dtype dtype is created.

  • num_krylov_vecs – The number of iterations (number of krylov vectors).

  • numeig – The nummber of eigenvector-eigenvalue pairs to be computed. If numeig > 1, reorthogonalize has to be True.

  • tol – The desired precision of the eigenvalus. Uses

  • which –

    [β€˜LM’ | β€˜SM’ | β€˜LR’ | β€˜SR’ | β€˜LI’] Which k eigenvectors and eigenvalues to find:

    ’LM’ : largest magnitude β€˜SM’ : smallest magnitude β€˜LR’ : largest real part β€˜SR’ : smallest real part β€˜LI’ : largest imaginary part

  • maxiter – The maximum number of iterations.

Returns

An array of numeig lowest eigenvalues list: A list of numeig lowest eigenvectors

Return type

np.ndarray

eigsh(A: Callable, args: Optional[List[Any]] = None, initial_state: Optional[Any] = None, shape: Optional[Tuple[int, ...]] = None, dtype: Optional[Type[numpy.number]] = None, num_krylov_vecs: int = 50, numeig: int = 1, tol: float = 1e-08, which: str = 'LR', maxiter: Optional[int] = None) β†’ Tuple[Any, List]#

Lanczos method for finding the lowest eigenvector-eigenvalue pairs of a symmetric (hermitian) linear operator A. A is a callable implementing the matrix-vector product. If no initial_state is provided then shape and dtype have to be passed so that a suitable initial state can be randomly generated. :param A: A (sparse) implementation of a linear operator :param arsg: A list of arguments to A. A will be called as

res = A(initial_state, *args).

Parameters
  • initial_state – An initial vector for the algorithm. If None, a random initial Tensor is created using the numpy.random.randn method.

  • shape – The shape of the input-dimension of A.

  • dtype – The dtype of the input A. If both no initial_state is provided, a random initial state with shape shape and dtype dtype is created.

  • num_krylov_vecs – The number of iterations (number of krylov vectors).

  • numeig – The nummber of eigenvector-eigenvalue pairs to be computed. If numeig > 1, reorthogonalize has to be True.

  • tol – The desired precision of the eigenvalus. Uses

  • which –

    [β€˜LM’ | β€˜SM’ | β€˜LR’ | β€˜SR’ | β€˜LI’ | β€˜SI’] Which k eigenvectors and eigenvalues to find:

    ’LM’ : largest magnitude β€˜SM’ : smallest magnitude β€˜LR’ : largest real part β€˜SR’ : smallest real part β€˜LI’ : largest imaginary part β€˜SI’ : smallest imaginary part

    Note that not all of those might be supported by specialized backends.

  • maxiter – The maximum number of iterations.

Returns

An array of numeig lowest eigenvalues list: A list of numeig lowest eigenvectors

Return type

Tensor

eigsh_lanczos(A: Callable, args: Optional[List[Any]] = None, initial_state: Optional[Any] = None, shape: Optional[Tuple] = None, dtype: Optional[Type[numpy.number]] = None, num_krylov_vecs: int = 20, numeig: int = 1, tol: float = 1e-08, delta: float = 1e-08, ndiag: int = 20, reorthogonalize: bool = False) β†’ Tuple[Any, List]#

Lanczos method for finding the lowest eigenvector-eigenvalue pairs of a linear operator A. :param A: A (sparse) implementation of a linear operator.

Call signature of A is res = A(vector, *args), where vector can be an arbitrary Tensor, and res.shape has to be vector.shape.

Parameters
  • arsg – A list of arguments to A. A will be called as res = A(initial_state, *args).

  • initial_state – An initial vector for the Lanczos algorithm. If None, a random initial Tensor is created using the backend.randn method

  • shape – The shape of the input-dimension of A.

  • dtype – The dtype of the input A. If both no initial_state is provided, a random initial state with shape shape and dtype dtype is created.

  • num_krylov_vecs – The number of iterations (number of krylov vectors).

  • numeig – The nummber of eigenvector-eigenvalue pairs to be computed. If numeig > 1, reorthogonalize has to be True.

  • tol – The desired precision of the eigenvalus. Uses np.linalg.norm(eigvalsnew[0:numeig] - eigvalsold[0:numeig]) < tol as stopping criterion between two diagonalization steps of the tridiagonal operator.

  • delta – Stopping criterion for Lanczos iteration. If a Krylov vector :math: x_n has an L2 norm \(\lVert x_n\rVert < delta\), the iteration is stopped. It means that an (approximate) invariant subspace has been found.

  • ndiag – The tridiagonal Operator is diagonalized every ndiag iterations to check convergence.

  • reorthogonalize – If True, Krylov vectors are kept orthogonal by explicit orthogonalization (more costly than reorthogonalize=False)

Returns

(eigvals, eigvecs)

eigvals: A list of numeig lowest eigenvalues eigvecs: A list of numeig lowest eigenvectors

eigvalsh(a: Any) β†’ Any[source]#

Get the eigenvalues of matrix a.

Parameters

a (Tensor) – tensor in matrix form

Returns

eigenvalues of a

Return type

Tensor

einsum(expression: str, *tensors: Any, optimize: bool = True) β†’ Any#

Calculate sum of products of tensors according to expression.

eps(dtype: Type[numpy.number]) β†’ float#

Return machine epsilon for given dtype

Parameters

dtype – A dtype.

Returns

Machine epsilon.

Return type

float

exp(tensor: Any) β†’ Any#

Return elementwise exp of tensor. :param tensor: A tensor.

Returns

Tensor

expm(a: Any) β†’ Any[source]#

Return expm log of matrix, matrix exponential. :param matrix: A tensor.

Returns

Tensor

eye(N: int, dtype: Optional[str] = None, M: Optional[int] = None) β†’ Any[source]#
Return an identity matrix of dimension dim

Depending on specific backends, dim has to be either an int (numpy, torch, tensorflow) or a ShapeType object (for block-sparse backends). Block-sparse behavior is currently not supported

Parameters
  • N (int) – The dimension of the returned matrix.

  • dtype – The dtype of the returned matrix.

  • M (int) – The dimension of the returned matrix.

from_dlpack(a: Any) β†’ Any#

Transform a dlpack capsule to a tensor

Parameters

a (Any) – the dlpack capsule

Returns

_description_

Return type

Tensor

gather1d(operand: Any, indices: Any) β†’ Any#

Return operand[indices], both operand and indices are rank-1 tensor.

Parameters
  • operand (Tensor) – rank-1 tensor

  • indices (Tensor) – rank-1 tensor with int dtype

Returns

operand[indices]

Return type

Tensor

get_random_state(seed: Optional[int] = None) β†’ Any#

Get the backend specific random state object.

Parameters

seed (Optional[int], optional) – [description], defaults to be None

:return:the backend specific random state object :rtype: Any

gmres(A_mv: Callable, b: Any, A_args: Optional[List] = None, A_kwargs: Optional[dict] = None, x0: Optional[Any] = None, tol: float = 1e-05, atol: Optional[float] = None, num_krylov_vectors: int = 20, maxiter: Optional[int] = 1, M: Optional[Callable] = None) β†’ Tuple[Any, int]#

GMRES solves the linear system A @ x = b for x given a vector b and a general (not necessarily symmetric/Hermitian) linear operator A.

As a Krylov method, GMRES does not require a concrete matrix representation of the n by n A, but only a function vector1 = A_mv(vector0, *A_args, **A_kwargs) prescribing a one-to-one linear map from vector0 to vector1 (that is, A must be square, and thus vector0 and vector1 the same size). If A is a dense matrix, or if it is a symmetric/Hermitian operator, a different linear solver will usually be preferable.

GMRES works by first constructing the Krylov basis K = (x0, A_mv@x0, A_mv@A_mv@x0, …, (A_mv^num_krylov_vectors)@x_0) and then solving a certain dense linear system K @ q0 = q1 from whose solution x can be approximated. For num_krylov_vectors = n the solution is provably exact in infinite precision, but the expense is cubic in num_krylov_vectors so one is typically interested in the num_krylov_vectors << n case. The solution can in this case be repeatedly improved, to a point, by restarting the Arnoldi iterations each time num_krylov_vectors is reached. Unfortunately the optimal parameter choices balancing expense and accuracy are difficult to predict in advance, so applying this function requires a degree of experimentation.

In a tensor network code one is typically interested in A_mv implementing some tensor contraction. This implementation thus allows b and x0 to be of whatever arbitrary, though identical, shape b = A_mv(x0, …) expects. Reshaping to and from a matrix problem is handled internally.

Parameters
  • A_mv – A function v0 = A_mv(v, *A_args, **A_kwargs) where v0 and v have the same shape.

  • b – The b in A @ x = b; it should be of the shape A_mv operates on.

  • A_args – Positional arguments to A_mv, supplied to this interface as a list. Default: None.

  • A_kwargs – Keyword arguments to A_mv, supplied to this interface as a dictionary. Default: None.

  • x0 – An optional guess solution. Zeros are used by default. If x0 is supplied, its shape and dtype must match those of b, or an error will be thrown. Default: zeros.

  • tol –

    Solution tolerance to achieve, norm(residual) <= max(tol*norm(b), atol). Default: tol=1E-05

    atol=tol

  • atol –

    Solution tolerance to achieve, norm(residual) <= max(tol*norm(b), atol). Default: tol=1E-05

    atol=tol

  • num_krylov_vectors –

    : Size of the Krylov space to build at each restart.

    Expense is cubic in this parameter. It must be positive. If greater than b.size, it will be set to b.size. Default: 20

  • maxiter – The Krylov space will be repeatedly rebuilt up to this many times. Large values of this argument should be used only with caution, since especially for nearly symmetric matrices and small num_krylov_vectors convergence might well freeze at a value significantly larger than tol. Default: 1.

  • M – Inverse of the preconditioner of A; see the docstring for scipy.sparse.linalg.gmres. This is only supported in the numpy backend. Supplying this argument to other backends will trigger NotImplementedError. Default: None.

Raises

ValueError – -if x0 is supplied but its shape differs from that of b. -in NumPy, if the ARPACK solver reports a breakdown (which usually indicates some kind of floating point issue). -if num_krylov_vectors is 0 or exceeds b.size. -if tol was negative. -if M was supplied with any backend but NumPy.

Returns

The converged solution. It has the same shape as b. info : 0 if convergence was achieved, the number of restarts otherwise.

Return type

x

grad(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) β†’ Callable[[...], Any][source]#

Return the function which is the grad function of input f.

Example

>>> f = lambda x,y: x**2+2*y
>>> g = tc.backend.grad(f)
>>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))
2
>>> g = tc.backend.grad(f, argnums=(0,1))
>>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))
[2, 2]
Parameters
  • f (Callable[..., Any]) – the function to be differentiated

  • argnums (Union[int, Sequence[int]], optional) – the position of args in f that are to be differentiated, defaults to be 0

Returns

the grad function of f with the same set of arguments as f

Return type

Callable[…, Any]

hessian(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0) β†’ Any#
i(dtype: Optional[Any] = None) β†’ Any[source]#

Return 1.j in as a tensor compatible with the backend.

Parameters

dtype (str) – β€œcomplex64” or β€œcomplex128”

Returns

1.j tensor

Return type

Tensor

imag(a: Any) β†’ Any[source]#

Return the elementwise imaginary value of a tensor a.

Parameters

a (Tensor) – tensor

Returns

imaginary value of a

Return type

Tensor

implicit_randc(a: Union[int, Sequence[int], Any], shape: Union[int, Sequence[int]], p: Optional[Union[Sequence[float], Any]] = None) β†’ Any#

[summary]

Parameters
  • g (Any) – [description]

  • a (Union[int, Sequence[int], Tensor]) – The possible options

  • shape (Union[int, Sequence[int]]) – Sampling output shape

  • p (Optional[Union[Sequence[float], Tensor]], optional) – probability for each option in a, defaults to None, as equal probability distribution

Returns

[description]

Return type

Tensor

implicit_randn(shape: Union[int, Sequence[int]] = 1, mean: float = 0, stddev: float = 1, dtype: str = '32') β†’ Any#

Call the random normal function with the random state management behind the scene.

Parameters
  • shape (Union[int, Sequence[int]], optional) – [description], defaults to 1

  • mean (float, optional) – [description], defaults to 0

  • stddev (float, optional) – [description], defaults to 1

  • dtype (str, optional) – [description], defaults to β€œ32”

Returns

[description]

Return type

Tensor

implicit_randu(shape: Union[int, Sequence[int]] = 1, low: float = 0, high: float = 1, dtype: str = '32') β†’ Any#

Call the random normal function with the random state management behind the scene.

Parameters
  • shape (Union[int, Sequence[int]], optional) – [description], defaults to 1

  • mean (float, optional) – [description], defaults to 0

  • stddev (float, optional) – [description], defaults to 1

  • dtype (str, optional) – [description], defaults to β€œ32”

Returns

[description]

Return type

Tensor

index_update(tensor: Any, mask: Any, assignee: Any) β†’ Any#

Update tensor at elements defined by mask with value assignee.

Parameters
  • tensor – A Tensor object.

  • mask – A boolean mask.

  • assignee – A scalar Tensor. The values to assigned to tensor at positions where mask is True.

inv(matrix: Any) β†’ Any#

Compute the matrix inverse of matrix.

Parameters

matrix – A matrix.

Returns

The inverse of matrix

Return type

Tensor

is_sparse(a: Any) β†’ bool[source]#

Determine whether the type of input a is sparse.

Parameters

a (Tensor) – input matrix a

Returns

a bool indicating whether the matrix a is sparse

Return type

bool

is_tensor(a: Any) β†’ bool[source]#

Return a boolean on whether a is a tensor in backend package.

Parameters

a (Tensor) – a tensor to be determined

Returns

whether a is a tensor

Return type

bool

item(tensor)#

Return the item of a 1-element tensor.

Parameters

tensor – A 1-element tensor

Returns

The value in tensor.

jacbwd(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0) β†’ Any#

Compute the Jacobian of f using reverse mode AD.

Parameters
  • f (Callable[..., Any]) – The function whose Jacobian is required

  • argnums (Union[int, Sequence[int]], optional) – the position of the arg as Jacobian input, defaults to 0

Returns

outer tuple for output, inner tuple for input args

Return type

Tensor

jacfwd(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0) β†’ Any#

Compute the Jacobian of f using the forward mode AD.

Parameters
  • f (Callable[..., Any]) – the function whose Jacobian is required

  • argnums (Union[int, Sequence[int]], optional) – the position of the arg as Jacobian input, defaults to 0

Returns

outer tuple for input args, inner tuple for outputs

Return type

Tensor

jacrev(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0) β†’ Any#

Compute the Jacobian of f using reverse mode AD.

Parameters
  • f (Callable[..., Any]) – The function whose Jacobian is required

  • argnums (Union[int, Sequence[int]], optional) – the position of the arg as Jacobian input, defaults to 0

Returns

outer tuple for output, inner tuple for input args

Return type

Tensor

jit(f: Callable[[...], Any], static_argnums: Optional[Union[int, Sequence[int]]] = None, jit_compile: Optional[bool] = None, **kws: Any) β†’ Callable[[...], Any][source]#

Return a jitted or graph-compiled version of fun for JAX backend. For all other backends returns fun. :param fun: Callable :param args: Arguments to fun. :param kwargs: Keyword arguments to fun.

Returns

jitted/graph-compiled version of fun, or just fun.

Return type

Callable

jvp(f: Callable[[...], Any], inputs: Union[Any, Sequence[Any]], v: Union[Any, Sequence[Any]]) β†’ Tuple[Union[Any, Sequence[Any]], Union[Any, Sequence[Any]]][source]#

Function that computes a (forward-mode) Jacobian-vector product of f. Strictly speaking, this function is value_and_jvp.

Parameters
  • f (Callable[..., Any]) – The function to compute jvp

  • inputs (Union[Tensor, Sequence[Tensor]]) – input for f

  • v (Union[Tensor, Sequence[Tensor]]) – tangents

Returns

(f(*inputs), jvp_tensor), where jvp_tensor is the same shape as the output of f

Return type

Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]

kron(a: Any, b: Any) β†’ Any[source]#

Return the kronecker product of two matrices a and b.

Parameters
  • a (Tensor) – tensor in matrix form

  • b (Tensor) – tensor in matrix form

Returns

kronecker product of a and b

Return type

Tensor

left_shift(x: Any, y: Any) β†’ Any[source]#

Shift the bits of an integer x to the left y bits.

Parameters
  • x (Tensor) – input values

  • y (Tensor) – Number of bits shift to x

Returns

result with the same shape as x

Return type

Tensor

log(tensor: Any) β†’ Any#

Return elementwise natural logarithm of tensor. :param tensor: A tensor.

Returns

Tensor

matmul(tensor1: Any, tensor2: Any) β†’ Any#

Perform a possibly batched matrix-matrix multiplication between tensor1 and tensor2. The following behaviour is similar to numpy.matmul: - If both arguments are 2-D they are multiplied like conventional

matrices.

  • If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly.

Both arguments to matmul have to be tensors of order >= 2. :param tensor1: An input tensor. :param tensor2: An input tensor.

Returns

The result of performing the matmul.

Return type

tensor

max(a: Any, axis: Optional[int] = None) β†’ Any[source]#

Return the maximum of an array or maximum along an axis.

Parameters
  • a (Tensor) – [description]

  • axis (Optional[int], optional) – [description], defaults to None

Returns

[description]

Return type

Tensor

mean(a: Any, axis: Optional[Sequence[int]] = None, keepdims: bool = False) β†’ Any[source]#

Compute the arithmetic mean for a along the specified axis.

Parameters
  • a (Tensor) – tensor to take average

  • axis (Optional[Sequence[int]], optional) – the axis to take mean, defaults to None indicating sum over flatten array

  • keepdims (bool, optional) – _description_, defaults to False

Returns

_description_

Return type

Tensor

min(a: Any, axis: Optional[int] = None) β†’ Any[source]#

Return the minimum of an array or minimum along an axis.

Parameters
  • a (Tensor) – [description]

  • axis (Optional[int], optional) – [description], defaults to None

Returns

[description]

Return type

Tensor

mod(x: Any, y: Any) β†’ Any[source]#

Compute y-mod of x (negative number behavior is not guaranteed to be consistent)

Parameters
  • x (Tensor) – input values

  • y (Tensor) – mod y

Returns

results

Return type

Tensor

multiply(tensor1: Any, tensor2: Any) β†’ Any#

Return the default multiplication of tensor.

A backend can override such implementation. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

Tensor

norm(tensor: Any) β†’ Any#

Calculate the L2-norm of the elements of tensor

numpy(a: Any) β†’ Any[source]#

Return the numpy array of a tensor a, but may not work in a jitted function.

Parameters

a (Tensor) – tensor in matrix form

Returns

numpy array of a

Return type

Tensor

one_hot(a: Any, num: int) β†’ Any#

See doc for onehot()

onehot(a: Any, num: int) β†’ Any[source]#

One-hot encodes the given a. Each index in the input a is encoded as a vector of zeros of length num with the element at index set to one:

Parameters
  • a (Tensor) – input tensor

  • num (int) – number of features in onehot dimension

Returns

onehot tensor with the last extra dimension

Return type

Tensor

ones(shape: Sequence[int], dtype: Optional[str] = None) β†’ Any[source]#

Return an ones-matrix of dimension dim Depending on specific backends, dim has to be either an int (numpy, torch, tensorflow) or a ShapeType object (for block-sparse backends). Block-sparse behavior is currently not supported :param shape: The dimension of the returned matrix. :type shape: int :param dtype: The dtype of the returned matrix.

outer_product(tensor1: Any, tensor2: Any) β†’ Any#

Calculate the outer product of the two given tensors.

pivot(tensor: Any, pivot_axis: int = - 1) β†’ Any#

Reshapes a tensor into a matrix, whose columns (rows) are the vectorized dimensions to the left (right) of pivot_axis.

In other words, with tensor.shape = (1, 2, 4, 5) and pivot_axis=2, this function returns an (8, 5) matrix.

Parameters
  • tensor – The tensor to pivot.

  • pivot_axis – The axis about which to pivot.

Returns

The pivoted tensor.

power(a: Any, b: Union[Any, float]) β†’ Any#
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.

Parameters
  • 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.

probability_sample(shots: int, p: Any, status: Optional[Any] = None, g: Optional[Any] = None) β†’ Any#

Drawn shots samples from probability distribution p, given the external randomness determined by uniform distributed status tensor or backend random generator g. This method is similar with stateful_randc, but it supports status beyond g, which is convenient when jit or vmap

Parameters
  • shots (int) – Number of samples to draw with replacement

  • p (Tensor) – prbability vector

  • status (Optional[Tensor], optional) – external randomness as a tensor with each element drawn uniformly from [0, 1], defaults to None

  • g (Any, optional) – backend random genrator, defaults to None

Returns

The drawn sample as an int tensor

Return type

Tensor

qr(tensor: Any, pivot_axis: int = - 1, non_negative_diagonal: bool = False) β†’ Tuple[Any, Any]#

Computes the QR decomposition of a tensor.

randn(shape: Tuple[int, ...], dtype: Optional[numpy.dtype] = None, seed: Optional[int] = None) β†’ Any#

Return a random-normal-matrix of dimension dim Depending on specific backends, dim has to be either an int (numpy, torch, tensorflow) or a ShapeType object (for block-sparse backends).

Block-sparse behavior is currently not supported :param shape: The dimension of the returned matrix. :type shape: int :param dtype: The dtype of the returned matrix. :param seed: The seed for the random number generator

random_split(key: Any) β†’ Tuple[Any, Any]#

A jax like split API, but it doesn’t split the key generator for other backends. It is just for a consistent interface of random code; make sure you know what the function actually does. This function is mainly a utility to write backend agnostic code instead of doing magic things.

Parameters

key (Any) – [description]

Returns

[description]

Return type

Tuple[Any, Any]

random_uniform(shape: Tuple[int, ...], boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), dtype: Optional[numpy.dtype] = None, seed: Optional[int] = None) β†’ Any#

Return a random uniform matrix of dimension dim.

Depending on specific backends, dim has to be either an int (numpy, torch, tensorflow) or a ShapeType object (for block-sparse backends). Block-sparse behavior is currently not supported :param shape: The dimension of the returned matrix. :type shape: int :param boundaries: The boundaries of the uniform distribution. :type boundaries: tuple :param dtype: The dtype of the returned matrix. :param seed: The seed for the random number generator

Returns

random uniform initialized tensor.

Return type

Tensor

real(a: Any) β†’ Any[source]#

Return the elementwise real value of a tensor a.

Parameters

a (Tensor) – tensor

Returns

real value of a

Return type

Tensor

relu(a: Any) β†’ Any[source]#

Rectified linear unit activation function. Computes the element-wise function:

\[\mathrm{relu}(x)=\max(x,0)\]
Parameters

a (Tensor) – Input tensor

Returns

Tensor after relu

Return type

Tensor

reshape(tensor: Any, shape: Any) β†’ Any#

Reshape tensor to the given shape.

Parameters

tensor – A tensor.

Returns

The reshaped tensor.

reshape2(a: Any) β†’ Any#

Reshape a tensor to the [2, 2, …] shape.

Parameters

a (Tensor) – Input tensor

Returns

the reshaped tensor

Return type

Tensor

reshapem(a: Any) β†’ Any#

Reshape a tensor to the [l, l] shape.

Parameters

a (Tensor) – Input tensor

Returns

the reshaped tensor

Return type

Tensor

reverse(a: Any) β†’ Any#

return a[::-1], only 1D tensor is guaranteed for consistent behavior

Parameters

a (Tensor) – 1D tensor

Returns

1D tensor in reverse order

Return type

Tensor

right_shift(x: Any, y: Any) β†’ Any[source]#

Shift the bits of an integer x to the right y bits.

Parameters
  • x (Tensor) – input values

  • y (Tensor) – Number of bits shift to x

Returns

result with the same shape as x

Return type

Tensor

rq(tensor: Any, pivot_axis: int = - 1, non_negative_diagonal: bool = False) β†’ Tuple[Any, Any]#

Computes the RQ (reversed QR) decomposition of a tensor.

scan(f: Callable[[Any, Any], Any], xs: Any, init: Any) β†’ Any#

This API follows tf.scan covention, i.e. no ys supported as jax

Parameters
  • f (Callable[Tuple[Tensor, Tensor], Tensor]) – _description_

  • xs (Tensor) – _description_

  • init (Tensor) – _description_

Returns

_description_

Return type

Tensor

scatter(operand: Any, indices: Any, updates: Any) β†’ Any[source]#

Roughly equivalent to operand[indices] = updates, indices only support shape with rank 2 for now.

Parameters
  • operand (Tensor) – [description]

  • indices (Tensor) – [description]

  • updates (Tensor) – [description]

Returns

[description]

Return type

Tensor

schur(a: Any, output: str = 'real') β†’ Tuple[Any, Any][source]#

Compute Schur decomposition of a matrix.

Parameters
  • a (Tensor) – _description_

  • output (str, optional) – _description_, defaults to β€œreal”

Returns

T, Z so that ZTZ^H = a

Return type

Tuple[Tensor, Tensor]

searchsorted(a: Any, v: Any, side: str = 'left') β†’ Any[source]#

Find indices where elements should be inserted to maintain order.

Parameters
  • a (Tensor) – input array sorted in ascending order

  • v (Tensor) – value to inserted

  • side (str, optional) – If β€˜left’, the index of the first suitable location found is given. If β€˜right’, return the last such index. If there is no suitable index, return either 0 or N (where N is the length of a), defaults to β€œleft”

Returns

Array of insertion points with the same shape as v, or an integer if v is a scalar.

Return type

Tensor

serialize_tensor(tensor: Any) β†’ str#

Return a string that serializes the given tensor.

Parameters

tensor – The input tensor.

Returns

A string representing the serialized tensor.

set_random_state(seed: Optional[int] = None, get_only: bool = False) β†’ Any[source]#

Set the random state attached to the backend.

Parameters
  • seed (Optional[int], optional) – the random seed, defaults to be None

  • get_only (bool, defaults to be False) – If set to be true, only get the random state in return instead of setting the state on the backend

shape_concat(values: Any, axis: int) β†’ Any#

Concatenate a sequence of tensors together about the given axis.

shape_prod(values: Any) β†’ Any#

Take the product of all of the elements in values

shape_tensor(tensor: Any) β†’ Any#

Get the shape of a tensor.

Parameters

tensor – A tensor.

Returns

The shape of the input tensor returned as another tensor.

shape_tuple(tensor: Any) β†’ Tuple[Optional[int], ...]#

Get the shape of a tensor as a tuple of integers.

Parameters

tensor – A tensor.

Returns

The shape of the input tensor returned as a tuple of ints.

sigmoid(a: Any) β†’ Any[source]#

Compute sigmoid of input a

Parameters

a (Tensor) – [description]

Returns

[description]

Return type

Tensor

sign(tensor: Any) β†’ Any#

Returns an elementwise tensor with entries y[i] = 1, 0, -1 tensor[i] > 0, == 0, and < 0 respectively.

For complex input the behaviour of this function may depend on the backend. The NumPy version returns y[i] = x[i]/sqrt(x[i]^2).

Parameters

tensor – The input tensor.

sin(a: Any) β†’ Any[source]#

Return sin of tensor. :param tensor: A tensor.

Returns

Tensor

sinh(a: Any) β†’ Any[source]#

Return the sinh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

sinh of a

Return type

Tensor

size(a: Any) β†’ Any[source]#

Return the total number of elements in a in tensor form.

Parameters

a (Tensor) – tensor

Returns

the total number of elements in a

Return type

Tensor

sizen(a: Any) β†’ int#

Return the total number of elements in tensor a, but in integer form.

Parameters

a (Tensor) – tensor

Returns

the total number of elements in tensor a

Return type

int

slice(tensor: Any, start_indices: Tuple[int, ...], slice_sizes: Tuple[int, ...]) β†’ Any#

Obtains a slice of a tensor based on start_indices and slice_sizes.

Parameters
  • 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.

softmax(a: Sequence[Any], axis: Optional[int] = None) β†’ Any[source]#

Softmax function. Computes the function which rescales elements to the range [0,1] such that the elements along axis sum to 1.

\[\mathrm{softmax}(x) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]
Parameters
  • a (Sequence[Tensor]) – Tensor

  • axis (int, optional) – A dimension along which Softmax will be computed , defaults to None for all axis sum.

Returns

concatenated tensor

Return type

Tensor

solve(A: Any, b: Any, assume_a: str = 'gen') β†’ Any[source]#

Solve the linear system Ax=b and return the solution x.

Parameters
  • A (Tensor) – The multiplied matrix.

  • b (Tensor) – The resulted matrix.

Returns

The solution of the linear system.

Return type

Tensor

sparse_dense_matmul(sp_a: Any, b: Any) β†’ Any[source]#

A sparse matrix multiplies a dense matrix.

Parameters
  • sp_a (Tensor) – a sparse matrix

  • b (Tensor) – a dense matrix

Returns

dense matrix

Return type

Tensor

sparse_shape(tensor: Any) β†’ Tuple[Optional[int], ...]#
sqrt(tensor: Any) β†’ Any#

Take the square root (element wise) of a given tensor.

sqrtmh(a: Any) β†’ Any#

Return the sqrtm of a Hermitian matrix a.

Parameters

a (Tensor) – tensor in matrix form

Returns

sqrtm of a

Return type

Tensor

stack(a: Sequence[Any], axis: int = 0) β†’ Any[source]#

Concatenates a sequence of tensors a along a new dimension axis.

Parameters
  • a (Sequence[Tensor]) – List of tensors in the same shape

  • axis (int, optional) – the stack axis, defaults to 0

Returns

concatenated tensor

Return type

Tensor

stateful_randc(g: numpy.random._generator.Generator, a: Union[int, Sequence[int], Any], shape: Union[int, Sequence[int]], p: Optional[Union[Sequence[float], Any]] = None) β†’ Any[source]#

[summary]

Parameters
  • g (Any) – [description]

  • a (Union[int, Sequence[int], Tensor]) – The possible options

  • shape (Union[int, Sequence[int]]) – Sampling output shape

  • p (Optional[Union[Sequence[float], Tensor]], optional) – probability for each option in a, defaults to None, as equal probability distribution

Returns

[description]

Return type

Tensor

stateful_randn(g: numpy.random._generator.Generator, shape: Union[int, Sequence[int]] = 1, mean: float = 0, stddev: float = 1, dtype: str = '32') β†’ Any[source]#

[summary]

Parameters
  • self (Any) – [description]

  • g (Any) – stateful register for each package

  • shape (Union[int, Sequence[int]]) – shape of output sampling tensor

  • mean (float, optional) – [description], defaults to 0

  • stddev (float, optional) – [description], defaults to 1

  • dtype (str, optional) – only real data type is supported, β€œ32” or β€œ64”, defaults to β€œ32”

Returns

[description]

Return type

Tensor

stateful_randu(g: numpy.random._generator.Generator, shape: Union[int, Sequence[int]] = 1, low: float = 0, high: float = 1, dtype: str = '32') β†’ Any[source]#

Uniform random sampler from low to high.

Parameters
  • g (Any) – stateful register for each package

  • shape (Union[int, Sequence[int]], optional) – shape of output sampling tensor, defaults to 1

  • low (float, optional) – [description], defaults to 0

  • high (float, optional) – [description], defaults to 1

  • dtype (str, optional) – only real data type is supported, β€œ32” or β€œ64”, defaults to β€œ32”

Returns

[description]

Return type

Tensor

std(a: Any, axis: Optional[Sequence[int]] = None, keepdims: bool = False) β†’ Any[source]#

Compute the standard deviation along the specified axis.

Parameters
  • a (Tensor) – _description_

  • axis (Optional[Sequence[int]], optional) – Axis or axes along which the standard deviation is computed, defaults to None, implying all axis

  • keepdims (bool, optional) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one, defaults to False

Returns

_description_

Return type

Tensor

stop_gradient(a: Any) β†’ Any[source]#

Stop backpropagation from a.

Parameters

a (Tensor) – [description]

Returns

[description]

Return type

Tensor

subtraction(tensor1: Any, tensor2: Any) β†’ Any#

Return the default substraction of tensor. A backend can override such implementation. :param tensor1: A tensor. :param tensor2: A tensor.

Returns

Tensor

sum(a: Any, axis: Optional[Sequence[int]] = None, keepdims: bool = False) β†’ Any#

Sum elements of tensor along the specified axis. Results in a new Tensor with the summed axis removed. :param tensor: An input tensor.

Returns

The result of performing the summation. The order of the tensor

will be reduced by 1.

Return type

tensor

svd(tensor: Any, pivot_axis: int = - 1, max_singular_values: Optional[int] = None, max_truncation_error: Optional[float] = None, relative: Optional[bool] = False) β†’ Tuple[Any, Any, Any, Any]#

Computes the singular value decomposition (SVD) of a tensor.

The SVD is performed by treating the tensor as a matrix, with an effective left (row) index resulting from combining the axes tensor.shape[:pivot_axis] and an effective right (column) index resulting from combining the axes tensor.shape[pivot_axis:].

For example, if tensor had a shape (2, 3, 4, 5) and pivot_axis was 2, then u would have shape (2, 3, 6), s would have shape (6), and vh would have shape (6, 4, 5).

If max_singular_values is set to an integer, the SVD is truncated to keep at most this many singular values.

If max_truncation_error > 0, as many singular values will be truncated as possible, so that the truncation error (the norm of discarded singular values) is at most max_truncation_error. If relative is set True then max_truncation_err is understood relative to the largest singular value.

If both max_singular_values and max_truncation_error are specified, the number of retained singular values will be min(max_singular_values, nsv_auto_trunc), where nsv_auto_trunc is the number of singular values that must be kept to maintain a truncation error smaller than max_truncation_error.

The output consists of three tensors u, s, vh such that: ```python

u[i1,…,iN, j] * s[j] * vh[j, k1,…,kM] == tensor[i1,…,iN, k1,…,kM]

``` Note that the output ordering matches numpy.linalg.svd rather than tf.svd.

Parameters
  • tensor – A tensor to be decomposed.

  • pivot_axis – Where to split the tensor’s axes before flattening into a matrix.

  • max_singular_values – The number of singular values to keep, or None to keep them all.

  • max_truncation_error – The maximum allowed truncation error or None to not do any truncation.

  • relative – Multiply max_truncation_err with the largest singular value.

Returns

Left tensor factor. s: Vector of ordered singular values from largest to smallest. vh: Right tensor factor. s_rest: Vector of discarded singular values (length zero if no

truncation).

Return type

u

switch(index: Any, branches: Sequence[Callable[[], Any]]) β†’ Any[source]#

branches[index]()

Parameters
  • index (Tensor) – [description]

  • branches (Sequence[Callable[[], Tensor]]) – [description]

Returns

[description]

Return type

Tensor

tan(a: Any) β†’ Any[source]#

Return the tan of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

tan of a

Return type

Tensor

tanh(a: Any) β†’ Any[source]#

Return the tanh of a tensor a.

Parameters

a (Tensor) – tensor in matrix form

Returns

tanh of a

Return type

Tensor

tensordot(a: Any, b: Any, axes: Union[int, Sequence[Sequence[int]]]) β†’ Any#

Do a tensordot of tensors a and b over the given axes.

Parameters
  • a – A tensor.

  • b – Another tensor.

  • axes – Two lists of integers. These values are the contraction axes.

tile(a: Any, rep: Any) β†’ Any[source]#

Constructs a tensor by tiling a given tensor.

Parameters
  • a (Tensor) – [description]

  • rep (Tensor) – 1d tensor with length the same as the rank of a

Returns

[description]

Return type

Tensor

to_dense(sp_a: Any) β†’ Any[source]#

Convert a sparse matrix to dense tensor.

Parameters

sp_a (Tensor) – a sparse matrix

Returns

the resulted dense matrix

Return type

Tensor

to_dlpack(a: Any) β†’ Any#

Transform the tensor a as a dlpack capsule

Parameters

a (Tensor) – _description_

Returns

_description_

Return type

Any

trace(tensor: Any, offset: int = 0, axis1: int = - 2, axis2: int = - 1) β†’ Any#

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.

Parameters
  • tensor – A tensor.

  • offset – Offset of the diagonal from the main diagonal.

  • axis1 – 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/last axis.

  • 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/last axis.

Returns

The batched summed diagonals.

Return type

array_of_diagonals

transpose(tensor: Any, perm: Optional[Sequence] = None) β†’ Any#

Transpose a tensor according to a given permutation. By default the axes are reversed. :param tensor: A tensor. :param perm: The permutation of the axes.

Returns

The transposed tensor

tree_flatten(pytree: Any) β†’ Tuple[Any, Any]#

Flatten python structure to 1D list

Parameters

pytree (Any) – python structure to be flattened

Returns

The 1D list of flattened structure and treedef which can be used for later unflatten

Return type

Tuple[Any, Any]

tree_map(f: Callable[[...], Any], *pytrees: Any) β†’ Any#

Return the new tree map with multiple arg function f through pytrees.

Parameters
  • f (Callable[..., Any]) – The function

  • pytrees (Any) – inputs as any python structure

Raises

NotImplementedError – raise when neither tensorflow or jax is installed.

Returns

The new tree map with the same structure but different values.

Return type

Any

tree_unflatten(treedef: Any, leaves: Any) β†’ Any#

Pack 1D list to pytree defined via treedef

Parameters
  • treedef (Any) – Def of pytree structure, the second return from tree_flatten

  • leaves (Any) – the 1D list of flattened data structure

Returns

Packed pytree

Return type

Any

unique_with_counts(a: Any, **kws: Any) β†’ Tuple[Any, Any][source]#

Find the unique elements and their corresponding counts of the given tensor a.

Parameters

a (Tensor) – [description]

Returns

Unique elements, corresponding counts

Return type

Tuple[Tensor, Tensor]

value_and_grad(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) β†’ Callable[[...], Tuple[Any, Any]][source]#

Return the function which returns the value and grad of f.

Example

>>> f = lambda x,y: x**2+2*y
>>> g = tc.backend.value_and_grad(f)
>>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))
5, 2
>>> g = tc.backend.value_and_grad(f, argnums=(0,1))
>>> g(tc.num_to_tensor(1),tc.num_to_tensor(2))
5, [2, 2]
Parameters
  • f (Callable[..., Any]) – the function to be differentiated

  • argnums (Union[int, Sequence[int]], optional) – the position of args in f that are to be differentiated, defaults to be 0

Returns

the value and grad function of f with the same set of arguments as f

Return type

Callable[…, Tuple[Any, Any]]

vectorized_value_and_grad(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0, vectorized_argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) β†’ Callable[[...], Tuple[Any, Any]][source]#

Return the VVAG function of f. The inputs for f is (args[0], args[1], args[2], …), and the output of f is a scalar. Suppose VVAG(f) is a function with inputs in the form (vargs[0], args[1], args[2], …), where vagrs[0] has one extra dimension than args[0] in the first axis and consistent with args[0] in shape for remaining dimensions, i.e. shape(vargs[0]) = [batch] + shape(args[0]). (We only cover cases where vectorized_argnums defaults to 0 here for demonstration). VVAG(f) returns a tuple as a value tensor with shape [batch, 1] and a gradient tuple with shape: ([batch]+shape(args[argnum]) for argnum in argnums). The gradient for argnums=k is defined as

\[g^k = \frac{\partial \sum_{i\in batch} f(vargs[0][i], args[1], ...)}{\partial args[k]}\]

Therefore, if argnums=0, the gradient is reduced to

\[g^0_i = \frac{\partial f(vargs[0][i])}{\partial vargs[0][i]}\]

, which is specifically suitable for batched VQE optimization, where args[0] is the circuit parameters.

And if argnums=1, the gradient is like

\[g^1_i = \frac{\partial \sum_j f(vargs[0][j], args[1])}{\partial args[1][i]}\]

, which is suitable for quantum machine learning scenarios, where f is the loss function, args[0] corresponds to the input data and args[1] corresponds to the weights in the QML model.

Parameters
  • f (Callable[..., Any]) – [description]

  • argnums (Union[int, Sequence[int]], optional) – [description], defaults to 0

  • vectorized_argnums (Union[int, Sequence[int]], defaults to 0) – the args to be vectorized, these arguments should share the same batch shape in the fist dimension

Returns

[description]

Return type

Callable[…, Tuple[Any, Any]]

vjp(f: Callable[[...], Any], inputs: Union[Any, Sequence[Any]], v: Union[Any, Sequence[Any]]) β†’ Tuple[Union[Any, Sequence[Any]], Union[Any, Sequence[Any]]][source]#

Function that computes the dot product between a vector v and the Jacobian of the given function at the point given by the inputs. (reverse mode AD relevant) Strictly speaking, this function is value_and_vjp.

Parameters
  • f (Callable[..., Any]) – the function to carry out vjp calculation

  • inputs (Union[Tensor, Sequence[Tensor]]) – input for f

  • v (Union[Tensor, Sequence[Tensor]]) – value vector or gradient from downstream in reverse mode AD the same shape as return of function f

Returns

(f(*inputs), vjp_tensor), where vjp_tensor is the same shape as inputs

Return type

Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]

vmap(f: Callable[[...], Any], vectorized_argnums: Union[int, Sequence[int]] = 0) β†’ Any[source]#

Return the vectorized map or batched version of f on the first extra axis. The general interface supports f with multiple arguments and broadcast in the fist dimension.

Parameters
  • f (Callable[..., Any]) – function to be broadcasted.

  • vectorized_argnums (Union[int, Sequence[int]], defaults to 0) – the args to be vectorized, these arguments should share the same batch shape in the fist dimension

Returns

vmap version of f

Return type

Any

vvag(f: Callable[[...], Any], argnums: Union[int, Sequence[int]] = 0, vectorized_argnums: Union[int, Sequence[int]] = 0, has_aux: bool = False) β†’ Callable[[...], Tuple[Any, Any]]#

Return the VVAG function of f. The inputs for f is (args[0], args[1], args[2], …), and the output of f is a scalar. Suppose VVAG(f) is a function with inputs in the form (vargs[0], args[1], args[2], …), where vagrs[0] has one extra dimension than args[0] in the first axis and consistent with args[0] in shape for remaining dimensions, i.e. shape(vargs[0]) = [batch] + shape(args[0]). (We only cover cases where vectorized_argnums defaults to 0 here for demonstration). VVAG(f) returns a tuple as a value tensor with shape [batch, 1] and a gradient tuple with shape: ([batch]+shape(args[argnum]) for argnum in argnums). The gradient for argnums=k is defined as

\[g^k = \frac{\partial \sum_{i\in batch} f(vargs[0][i], args[1], ...)}{\partial args[k]}\]

Therefore, if argnums=0, the gradient is reduced to

\[g^0_i = \frac{\partial f(vargs[0][i])}{\partial vargs[0][i]}\]

, which is specifically suitable for batched VQE optimization, where args[0] is the circuit parameters.

And if argnums=1, the gradient is like

\[g^1_i = \frac{\partial \sum_j f(vargs[0][j], args[1])}{\partial args[1][i]}\]

, which is suitable for quantum machine learning scenarios, where f is the loss function, args[0] corresponds to the input data and args[1] corresponds to the weights in the QML model.

Parameters
  • f (Callable[..., Any]) – [description]

  • argnums (Union[int, Sequence[int]], optional) – [description], defaults to 0

  • vectorized_argnums (Union[int, Sequence[int]], defaults to 0) – the args to be vectorized, these arguments should share the same batch shape in the fist dimension

Returns

[description]

Return type

Callable[…, Tuple[Any, Any]]

zeros(shape: Sequence[int], dtype: Optional[str] = None) β†’ Any[source]#

Return a zeros-matrix of dimension dim Depending on specific backends, dim has to be either an int (numpy, torch, tensorflow) or a ShapeType object (for block-sparse backends).

Block-sparse behavior is currently not supported :param shape: The dimension of the returned matrix. :type shape: int :param dtype: The dtype of the returned matrix.