This package contains unified functions for framework-agnostic tensor and array manipulation. Given the input tensor-like object, the call is dispatched to the corresponding array manipulation framework, allowing for end-to-end differentiation to be preserved.


These functions are experimental, and only a subset of common functionality is supported. Furthermore, the names and behaviour of these functions may differ from similar functions in common frameworks; please refer to the function docstrings for more details.

The following frameworks are currently supported:

  • NumPy

  • Autograd

  • TensorFlow

  • PyTorch

  • JAX



Determines the correct framework to dispatch to given a sequence of tensor-like objects.


Combine a sequence of 2D tensors to form a block diagonal tensor.

concatenate(values[, axis])

Concatenate a sequence of tensors along the specified axis.

diag(values[, k])

Construct a diagonal tensor from a list of scalars.

dot(tensor1, tensor2)

Returns the matrix or dot product of two tensors.

ones_like(tensor[, dtype])

Returns a tensor of all ones with the same shape and dtype as the input tensor.

stack(values[, axis])

Stack a sequence of tensors along the specified axis.

where(condition, x, y)

Returns elements chosen from x or y depending on a boolean tensor condition.

allclose(a, b[, rtol, atol])

Returns True if two arrays are element-wise equal within a tolerance.

allequal(tensor1, tensor2, **kwargs)

Returns True if two tensors are element-wise equal along a given axis.

cast(tensor, dtype)

Casts the given tensor to a new type.

cast_like(tensor1, tensor2)

Casts a tensor to the same dtype as another.

convert_like(tensor1, tensor2)

Convert a tensor to the same type as another.


Returns the name of the package that any array/tensor manipulations will dispatch to.

requires_grad(tensor[, interface])

Returns True if the tensor is considered trainable.

cov_matrix(prob, obs[, wires, diag_approx])

Calculate the covariance matrix of a list of commuting observables, given the joint probability distribution of the system in the shared eigenbasis.

marginal_prob(prob, axis)

Compute the marginal probability given a joint probability distribution expressed as a tensor.

unwrap(values[, max_depth])

Unwrap a sequence of objects to NumPy arrays.

frobenius_inner_product(A, B[, normalize])

Frobenius inner product between two matrices.


Returns a set containing the trainable indices of a sequence of values.

is_independent(func, interface, args[, …])

Test whether a function is independent of its input arguments, both numerically and analytically.