pytaco.matmul

pytaco.matmul(t1, t2, out_format=mode_format(compressed), dtype=None)

Matrix multiplication of two tensors.

Parameters
t1, t2: tensors, array_like

The tensors to multiply. Tensors must be at least 2-D.

out_format: format, mode_format, optional

  • If a format is specified, the result tensor is stored in the format out_format.

  • If a mode_format is specified, the result the result tensor has a with all of the dimensions stored in the mode_format passed in.

dtype: Datatype

The datatype of the output tensor.

Returns
res: tensor

The matrix product of the input tensors.

Notes

  • When both arguments are 2-D conventional matrix multiply is performed.

  • If either argument is N-D N > 2, the tensor is treated as a stack of matrices and is broadcasted in this manner.

Examples

Here we demonstrate broadcasting for a 3-D tensor. We use NumPy arrays for demonstration due to easy data generation. However, we could have given sparse tensors of any format for taco to compute. Note that the choice of a tensor format has a big effect on the final performance. For instance it is favorable to multiply CSR matrices with CSC since dot products can be easily computed.

>>> import pytaco as pt
>>> import numpy as np
>>> a = np.arange(2 * 2 * 4).reshape((2, 2, 4))
>>> b = np.arange(2 * 2 * 4).reshape((2, 4, 2))
>>> pt.matmul(a,b).shape
[2, 2, 2]
>>> pt.matmul(a, b)[0, 1, 1]
98