sponge.function.gather_vector

sponge.function.gather_vector(tensor: Tensor, index: Tensor)[source]

Gather vector from the penultimate axis (axis=-2) of the tensor according to index.

Parameters

tensor (Tensor) – Tensor of shape \((B, X, D)\), where \(B\) is batch size, \(X\) is an arbitrary value., \(D\) is spatial dimension of the simulation system, usually is 3.
index (Tensor) – Tensor of shape \((B, ...,)\). Data type is int.

Returns

vector (Tensor), a tensor of shape \((B, ..., D)\)