fit_a_nef.tasks.shape.stack_cut_tensor

fit_a_nef.tasks.shape.stack_cut_tensor(tensor_list: Sequence[Array], repeat_upto: int = 100000, shuffle: bool = False)

Repeat each tensor by repeat and stack them over batch dimension. Tensor is trimmed to minimum size of the tensor list.

Args:

tensor_list (Sequence[jnp.ndarray]): The list of tensors to concatenate and pad. repeat_upto (int, optional): Upper limit of times to repeat each tensor. Defaults to 1. shuffle (bool, optional): Whether to shuffle the tensors in the repeating step. Defaults to False.

Returns:

jnp.ndarray: The concatenated and cut tensor.