einx.argmax

Contents

einx.argmax#

einx.argmax(description, tensor, *, backend=None, **parameters)[source]#

Find the coordinates of the maximum values in the given tensor.

The elementary operation has the signature [...] -> [n]. It takes an n-dimensional tensor as input and returns the n-dimensional coordinate vector of the maximum value.

For 1-dimensional tensors, the elementary operation also accepts the signature [a] -> []. For example, the following two operations compute the same output:

y = einx.argmaargmaxx("a [b] -> a [1]", x)
y = einx.argmax("a [b] -> a ", x)

If no output is given, it is determined implicitly by replacing a single bracketed expression in the input with [n]. For example, the following two operations compute the same output:

y = einx.argmax("a [b c]", x)
y = einx.argmax("a [b c] -> a [2]", x)
Parameters:

  • description (str) – Description string for the operation in einx notation.

  • tensor (Tensor) – Input tensor or tensor factory matching the description string.

  • backend (Union[Backend, str, None]) – Backend to use for all operations. If None, uses the default backend for the given setting. Defaults to None.

  • graph – Whether to return the compiled code representation of this operation instead of computing the result. Defaults to False.

  • **parameters (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[bool | int | float | complex | str | bytes]]) – Additional parameters that specify dimension sizes, e.g. a=4.

Return type:

Tensor

Returns:

The result of the operation if graph=False, otherwise the compiled code representation of the operation.

Contents