jax.numpy.einsum#
- jax.numpy.einsum(subscripts, /, *operands, out=None, optimize='optimal', precision=None, preferred_element_type=None, _use_xeinsum=False, _dot_general=<function dot_general>)[source]#
Evaluates the Einstein summation convention on the operands.
LAX-backend implementation of
numpy.einsum()
.In addition to the original NumPy arguments listed below, also supports
precision
for extra control over matrix-multiplication precision on supported devices.precision
may be set toNone
, which means default precision for the backend, aPrecision
enum value (Precision.DEFAULT
,Precision.HIGH
orPrecision.HIGHEST
) or a tuple of twoPrecision
enums indicating separate precision for each argument. A tupleprecision
does not necessarily map to multiple arguments ofeinsum()
; rather, the specifiedprecision
is forwarded to eachdot_general
call used in the implementation.jax.numpy.einsum()
also differs fromnumpy.einsum()
in that theoptimize
keyword defaults to"optimal"
rather thanFalse
.Original docstring below.
Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be represented in a simple fashion. In implicit mode einsum computes these values.
In explicit mode, einsum provides further flexibility to compute other array operations that might not be considered classical Einstein summation operations, by disabling, or forcing summation over specified subscript labels.
See the notes and examples for clarification.
- 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.
optimize ({False, True, 'greedy', 'optimal'}, optional) – Controls if intermediate optimization should occur. No optimization will occur if False and True will default to the ‘greedy’ algorithm. Also accepts an explicit contraction list from the
np.einsum_path
function. Seenp.einsum_path
for more details. Defaults to False.out (None) –
precision (PrecisionLike) –
preferred_element_type (DTypeLike | None) –
_use_xeinsum (bool) –
_dot_general (Callable[..., Array]) –
- Returns:
output – The calculation based on the Einstein summation convention.
- Return type:
ndarray