TensorArray, TensorArray. This is mostly intended for loops, see for example scan().

class returnn.frontend.tensor_array.TensorArray(tensor_template: Tensor, *, _backend_tensor_array: Any | None = None, _backend: Type[Backend] | None = None, _enable_delayed_check: bool = False)[source]#

TensorArray. Think of this like a list of tensors. E.g. if each tensor has shape (B, D), and we have N tensors, stacking them together would give us a tensor of shape (N, B, D). Reversely, unstacking a tensor of shape (N, B, D) on the N axis would give us a list of N tensors of shape (B, D).

We use a functional API, and each modifying operation (push_back) returns a new TensorArray object. This is to make sure it works well together with both eager-based and graph-based frameworks.

Internally, the backend functions give us some opaque tensor array object (e.g. TF TensorArray, or maybe just a pure Python list of tensors in case of eager-based frameworks).

classmethod unstack(tensor: Tensor, *, axis: Dim) TensorArray[source]#
push_back(tensor: Tensor) TensorArray[source]#
stack(*, axis: Dim) Tensor[source]#