GatherMM

Matrix multiplication with matrix-level gather.

Abstract Signature:

GatherMM(a: Tensor, b: Tensor, lhs_indices: Tensor, rhs_indices: Tensor)

PyTorchJAX (Core)Apple MLX

PyTorch

API: —

Strategy: Custom / Partial

JAX (Core)

API: —

Strategy: Custom / Partial

Apple MLX

API: mlx.core.gather_mm

Strategy: Direct Mapping

Official Docs ↗