einx.add_at

Contents

einx.add_at#

einx.add_at(description, *tensors, backend=None, **parameters)[source]#

Adds values from an update tensor to a target tensor at the coordinates specified by an indexing tensor.

The elementary operation has the signature [...] , [n], [] -> [...]. The first argument is the n-dimensional value tensor, the second argument specifies a single n-dimensional coordinate, the third argument is the added value, and the result is the value tensor with the value added at that coordinate.

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

y = einx.add_at("p [h], p [1], p -> p [h]", x, idx,       update)
y = einx.add_at("p [h], p,     p -> p [h]", x, idx[:, 0], update)

The elementary operation also accepts multiple coordinate tensors as input, in which case they are concatenated first. The length of the resulting coordinate vector must equal the number of dimensions of the value tensor. The update tensor always is the last argument. For example, the following two operations compute the same output:

y = einx.add_at("p [a b c d], p [n],           p -> p", x, idx,                                 update)
y = einx.add_at("p [a b c d], p, p [2], p [1], p -> p", x, idx[:, 0], idx[:, 1:3], idx[:, 3:4], update)

If no output expression is given, it is implicitly chosen to be the same as the input expression of the value tensor. For example, the following two operations compute the same output:

y = einx.add_at("b [h w] c, b p [2], b p c", x, idx, update)
y = einx.add_at("b [h w] c, b p [2], b p c -> b [h w] c", x, idx, update)

The order in which the updates are applied depends on the chosen backend. The operation also may or may not update the target tensor inplace. Please check by inspecting the code representation (by passing graph=True).

Parameters:

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

  • *tensors (Tensor) – Input tensors or tensor factories 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