tensordot

Contents

tensordot#

class brainunit.math.tensordot(a, b, axes=2, precision=None, preferred_element_type=None, **kwargs)#

Compute tensor dot product along specified axes.

The resulting unit is a.unit * b.unit.

Parameters:

  • a (Array | ndarray | bool | number | bool | int | float | complex | saiunit.Quantity) – First tensor.

  • b (Array | ndarray | bool | number | bool | int | float | complex | saiunit.Quantity) – Second tensor.

  • axes (int | Sequence[int] | Sequence[Sequence[int]]) – If an int N, sum over the last N axes of a and the first N axes of b. Or a list of two sequences of axis indices.

  • precision (Any) – Either None (default) or a Precision enum value, or a tuple of two such values.

  • preferred_element_type (str | type[Any] | dtype | SupportsDType | None) – Accumulation and result dtype.

Returns:

output – The tensor dot product. The resulting unit is a.unit * b.unit.

Return type:

Array | saiunit.Quantity

Examples

>>> import saiunit as u
>>> a = u.math.array([[1.0, 2.0], [3.0, 4.0]]) * u.meter
>>> b = u.math.array([[5.0, 6.0], [7.0, 8.0]]) * u.second
>>> u.math.tensordot(a, b, axes=1)  # unit is meter * second
Contents