pytblis package#

pytblis.ascontiguousarray(a)[source]#

Parallel transpose the input to C-contiguous layout.

Parameters:

a (array_like) – Input array.

Returns:

Contiguous array.

Return type:

ndarray

pytblis.contract(subscripts, a, b, alpha=1.0, beta=0.0, out=None, conja=False, conjb=False, allow_partial_trace=False)[source]#

Perform tensor contraction based on the provided subscripts. C (stored in out if provided) is computed as: C = alpha * einsum(subscripts, a, b) + beta * C if out is provided.

Parameters:

  • subscripts (str) – Subscripts defining the contraction.

  • a (array_like) – First tensor operand.

  • b (array_like) – Second tensor operand.

  • alpha (scalar, optional) – Scaling factor for the product of a and b.

  • beta (scalar, optional) – Scaling factor for the output tensor. Must be 0.0 if out is None.

  • conja (bool, optional) – If True, conjugate the first tensor a before contraction. Alpha is not conjugated.

  • conjb (bool, optional) – If True, conjugate the second tensor b before contraction. Beta is not conjugated.

  • allow_partial_trace (bool, optional) – If True, handle redundant indices in subscripts for a and b by doing partial trace before contraction.

  • out (array_like, optional) – Output tensor to store the result.

Returns:

Result of the tensor contraction.

Return type:

ndarray

pytblis.einsum(*operands, out=None, optimize='greedy', **kwargs)[source]#
einsum(subscripts, *operands, out=None, order=’K’,

optimize=’greedy’)

Evaluates the Einstein summation convention on the operands.

Drop-in replacement for numpy.einsum, using TBLIS for tensor contractions.

Parameters:

  • subscripts (str) – Specifies the subscripts for summation as comma separated list of subscript labels. An implicit (classical Einstein summation) calculation is performed unless the explicit indicator ‘->’ is included as well as subscript labels of the precise output form.

  • operands (list of array_like) – These are the arrays for the operation.

  • out (ndarray, optional) – If provided, the calculation is done into this array.

  • order ({'C', 'F', 'A', 'K'}, optional) – Controls the memory layout of the output. ‘C’ means it should be C contiguous. ‘F’ means it should be Fortran contiguous, ‘A’ means it should be ‘F’ if the inputs are all ‘F’, ‘C’ otherwise. ‘K’ means it should be as close to the layout as the inputs as is possible, including arbitrarily permuted axes. Default is ‘K’.

  • optimize ({'greedy', 'optimal'}, default 'greedy') – Controls the optimization strategy used to compute the contraction.

Returns:

output – The calculation based on the Einstein summation convention.

Return type:

ndarray

class pytblis.reduce_t(*values)#

Bases: Enum

REDUCE_MAX = 2#
REDUCE_MAX_ABS = 3#
REDUCE_MIN = 4#
REDUCE_MIN_ABS = 5#
REDUCE_NORM_1 = 1#
REDUCE_NORM_2 = 6#
REDUCE_NORM_INF = 3#
REDUCE_SUM = 0#
REDUCE_SUM_ABS = 1#
pytblis.tensordot(a, b, axes=2)[source]#

Compute tensor dot product along specified axes using TBLIS.

Given two tensors, a and b, and an array_like object containing two array_like objects, (a_axes, b_axes), sum the products of a’s and b’s elements (components) over the axes specified by a_axes and b_axes. The third argument can be a single non-negative integer_like scalar, N; if it is such, then the last N dimensions of a and the first N dimensions of b are summed over.

Parameters:

  • a (array_like) – Tensors to “dot”.

  • b (array_like) – Tensors to “dot”.

  • axes (int or (2,) array_like) –

    • integer_like If an int N, sum over the last N axes of a and the first N axes of b in order. The sizes of the corresponding axes must match.

    • (2,) array_like Or, a list of axes to be summed over, first sequence applying to a, second to b. Both elements array_like must be of the same length.

Returns:

output – The tensor dot product of the input.

Return type:

ndarray

pytblis.transpose_add(subscripts, a, alpha=1.0, beta=0.0, out=None, conja=False, conjout=False)[source]#

Perform tensor transpose and addition based on the provided subscripts. High-level wrapper for tblis_tensor_add. B (stored in out if provided) is computed as: B = alpha * transpose(a, subscripts) + beta * B. Optionally, conjugation can be applied to a and out.

Parameters:

  • subscripts (str) – Subscripts defining the contraction.

  • a (array_like) – Tensor operand

  • alpha (scalar, optional) – Scaling factor for a.

  • beta (scalar, optional) – Scaling factor for the output tensor b. Must be 0.0 if out is None.

  • conja (bool, optional) – If True, conjugate the first tensor a. Alpha is not conjugated.

  • conjout (bool, optional) – If True, conjugate the output tensor b. Beta is not conjugated.

  • out (array_like, optional) – Output tensor containing b.

Returns:

Result of the tensor contraction.

Return type:

ndarray

Examples

>>> import numpy as np
>>> from pytblis import transpose_add
>>> a = np.random.rand(3, 4, 5)
>>> b = transpose_add("ijk->ikj", a, alpha=1.0)
pytblis.add(A: ndarray[], B: ndarray[], idx_A: str, idx_B: str, alpha: complex = 1.0, beta: complex = 1.0, conja: bool = False, conjb: bool = False) → None#

Performs the operation B[idx_B] := alpha * A[idx_A] + beta * B[idx_B]. Equivalent to the einsum operation B = alpha * einsum(f’{idx_A}->{idx_B}’, A) + beta * B. The tensors A and B must have the same type and compatible shapes.

Parameters:

  • A (ndarray) – Tensor A (not overwritten).

  • B (ndarray) – Tensor B (overwritten).

  • idx_A (str) – Indices for tensor A.

  • idx_B (str) – Indices for tensor B.

  • alpha (scalar, optional) – Scalar multiplier for A (default is 1.0).

  • beta (scalar, optional) – Scalar multiplier for B (default is 1.0).

  • conja (bool, optional) – If True, conjugate tensor A (default is False); alpha is not conjugated.

  • conjb (bool, optional) – If True, conjugate tensor B (default is False); beta is not conjugated.

Return type:

None

pytblis.mult(A: ndarray[], B: ndarray[], C: ndarray[], idx_A: str, idx_B: str, idx_C: str, alpha: complex = 1.0, beta: complex = 0.0, conja: bool = False, conjb: bool = False) → None#

Generalized tensor multiplication. Computes C[idx_C] = alpha * A[idx_A] * B[idx_B] + beta * C[idx_C]. Equivalent to the einsum operation C = alpha * einsum(f’{idx_A},{idx_B}->{idx_C}’, A, B) + beta * C. The tensors A, B, and C must have the same type and compatible shapes.

Parameters:

  • A (ndarray) – Tensor A (not overwritten).

  • B (ndarray) – Tensor B (not overwritten).

  • C (ndarray) – Tensor C (overwritten).

  • idx_A (str) – Indices for tensor A.

  • idx_B (str) – Indices for tensor B.

  • idx_C (str) – Indices for tensor C.

  • alpha (scalar, optional) – Scalar multiplier for A and B (default is 1.0).

  • beta (scalar, optional) – Scalar multiplier for C (default is 0.0).

  • conja (bool, optional) – If True, conjugate tensor A (default is False); alpha is not conjugated.

  • conjb (bool, optional) – If True, conjugate tensor B (default is False); beta is not conjugated.

Return type:

None

pytblis.dot(A: ndarray[], B: ndarray[], idx_A: str, idx_B: str, alpha: complex = 1.0, beta: complex = 1.0, conja: bool = False, conjb: bool = False) → object#

Computes the dot product of tensors A and B. The result is a scalar value, equal to alpha * beta * einsum(f’{idx_A},{idx_B}’)->’’.

Parameters:

  • A (ndarray) – Tensor A (not overwritten).

  • B (ndarray) – Tensor B (not overwritten).

  • idx_A (str) – Indices for tensor A.

  • idx_B (str) – Indices for tensor B.

  • alpha (scalar, optional) – Scalar multiplier for A (default is 1.0).

  • beta (scalar, optional) – Scalar multiplier for B (default is 1.0).

  • conja (bool, optional) – If True, conjugate tensor A (default is False); alpha is not conjugated.

  • conjb (bool, optional) – If True, conjugate tensor B (default is False); beta is not conjugated.

Returns:

alpha * A[idx_A] dot beta * B[idx_B]

Return type:

scalar

pytblis.reduce(A: ndarray[], idx_A: str, op: pytblis._pytblis_impl.reduce_t, conja: bool = False) → object#

Tensor reduction operation. Computes the reduction of tensor A over indices idx_A using the operation op. The result is a scalar value or a tuple (scalar, index) if the operation returns an index.

Parameters:

  • A (ndarray) – Tensor A (not overwritten).

  • idx_A (str) – Indices for tensor A.

  • op (pytblis.reduce_t) – Reduction operation to perform.

  • conja (bool, optional) – If True, conjugate tensor A (default is False); the scalar is not conjugated.

Returns:

The result of the reduction operation. If the operation returns an index, a tuple (scalar, index) is returned, where scalar is the result of the reduction operation and index is the index of the maximum or minimum value in the tensor A. Otherwise, just the scalar result is returned.

Return type:

scalar or tuple

pytblis.shift(A: ndarray[], idx_A: str, alpha: complex = 0.0, beta: complex = 1.0) → None#

A_[idx_A] = alpha + beta * A_[idx_A]

pytblis.get_num_threads() → int#

Get the number of threads used by TBLIS.

pytblis.set_num_threads(num_threads: int) → None#

Set the number of threads used by TBLIS.