tensorplay.func.rearrange

tensorplay.func.rearrange(tensor: Any, pattern: str, **axes_lengths: int) → Any[source]

Reshapes and permutes tensor according to pattern.

Parameters:

• tensor – the tensor to rearrange, or a sequence of tensors, which is stacked along a new leading axis first.

• pattern (str) – "<input axes> -> <output axes>". Names bind by position on the left and select by name on the right. Parentheses group axes; ... matches any number of axes.

• **axes_lengths – sizes for axes the pattern splits an input axis into, which cannot be inferred from the shape alone.

Example

>>> x = tensorplay.randn(2, 3, 4)
>>> rearrange(x, "b h w -> b w h").shape
tensorplay.Size(2, 4, 3)
>>> rearrange(x, "b h w -> b (h w)").shape
tensorplay.Size(2, 12)
>>> rearrange(x, "(b1 b2) h w -> b1 b2 h w", b1=1).shape
tensorplay.Size(1, 2, 3, 4)