Unflatten

Expands a dimension of the input tensor over multiple dimensions.

Abstract Signature:

Unflatten(input: Tensor, dim: int, sizes: Tuple[int])

PyTorch

API: torch.unflatten
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: —
Strategy: Plugin (unflatten_logic)

NumPy

API: —
Strategy: Plugin (unflatten_manual)

Keras

API: keras.ops.reshape
Strategy: Plugin (unflatten_logic)

TensorFlow

API: tf.reshape
Strategy: Plugin (unflatten_manual)

Apple MLX

API: mlx.core.unflatten
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: nnx.graph.MergeContext.unflatten
Strategy: Plugin (unflatten_manual)