enp.flatten

etils.enp.flatten(

array: etils.enp.einops_utils._ArrayT,

pattern: str,

) → tuple[etils.enp.einops_utils._ArrayT, tuple[int, ...]]

Flatten an array along custom dimensions.

Uses einops syntax.

flat_x, batch_shape = enp.flatten(x, '... h w c')
y = enp.unflatten(y, batch_shape, '... h w c')

• x.shape == (h, w, c) -> flat_x.shape == (1, h, w, c)

• x.shape == (b, h, w, c) -> flat_x.shape == (b, h, w, c)

• x.shape == (b, n, h, w, c) -> flat_x.shape == (b * n, h, w, c)

Parameters:

• array – Array to flatten.

• pattern – Einops pattern to flatten the array.

Returns:

Tuple of (flattened array, batch shape).