tf_agents.utils.common.transpose_batch_time

View source on GitHub

Transposes the batch and time dimensions of a Tensor.

If the input tensor has rank < 2 it returns the original tensor. Retains as much of the static shape information as possible.

x A Tensor.

x transposed along the first two dimensions.