torch.nn.functional.grouped_mm

mat_a (Tensor) – Left operand. When 2D, its leading dimension is sliced into groups according to offs. When 3D, its first dimension enumerates the groups directly and offs must be None.