NativeChannelShuffle

Native kernel level implementation of the channel_shuffle.

Abstract Signature:

NativeChannelShuffle(input: Tensor, groups: int)

PyTorchJAX (Core)

PyTorch

API: torch.native_channel_shuffle

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.numpy.reshape

Strategy: Plugin (channel_shuffle_macro)