(l-onnx-doc-GroupedMatMul)=

# GroupedMatMul


(l-onnx-op-groupedmatmul-29)=

## GroupedMatMul - 29

### Version

- **name**: [GroupedMatMul (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#GroupedMatMul)
- **domain**: `main`
- **since_version**: `29`
- **function**: `True`
- **support_level**: `SupportType.COMMON`
- **shape inference**: `True`

This version of the operator has been available
**since version 29**.

### Summary

GroupedMatMul multiplies each row of a token matrix by one or more selected
group (expert) weight matrices. Given `input` of shape `[M, K]`, `weights` of
shape `[G, K, N]`, and `group_indices` of shape `[M, S]`, where `S` is the number
of selections per token, the output has shape `[M, S, N]` and is defined by:

```
output[m, i] = input[m] @ weights[group_indices[m, i]]
```

If `bias` is present, `bias[group_indices[m, i]]` is added to each corresponding
result. Every value in `group_indices` must be in the range `[0, G)`.
Values outside this range, including negative indices, are invalid input;
the behavior is undefined for such inputs. In particular, negative indices
do not select groups by counting backwards from `G`.

This operator is represented as a context-dependent function, which also provides
its reference implementation. Its decomposition uses Gather, Expand, and MatMul
to define the result, while runtimes are expected to use a fused
grouped-matrix-multiplication implementation to avoid materializing the expanded
intermediate tensors, including selected weights of shape `[M * S, K, N]`.

Weight-only quantized experts are intended to be expressed by applying
DequantizeLinear to quantized weights (with the appropriate scales and zero
points) and passing the result to GroupedMatMul. Runtimes may fuse this
composition to avoid materializing the dequantized `[G, K, N]` weights.

### Inputs

Between 3 and 4 inputs.

- **input** (heterogeneous) - **T**:

  Token matrix with shape [M, K].
- **weights** (heterogeneous) - **T**:

  Stack of G group weight matrices with shape [G, K, N].
- **group_indices** (heterogeneous) - **Tind**:

  Group index for each token and selection, with shape [M, S]. Values outside [0, G) are invalid input.
- **bias** (optional, heterogeneous) - **T**:

  Optional per-group bias with shape [G, N].

### Outputs

- **output** (heterogeneous) - **T**:

  Per-group matrix multiplication results with shape [M, S, N].

### Type Constraints

* **T** in ( `tensor(bfloat16)`, `tensor(double)`, `tensor(float)`, `tensor(float16)` ):

  Constrain input, weights, bias, and output to floating-point tensors.
* **Tind** in ( `tensor(int32)`, `tensor(int64)` ):

  Constrain group indices to int32 or int64 tensors.

### Examples

#### _additional_types

```python
import numpy as np
import onnx

for dtype in (np.float16, np.float64):
    for index_dtype in (np.int32, np.int64):
        node = onnx.helper.make_node(
            "GroupedMatMul",
            inputs=["input", "weights", "group_indices", "bias"],
            outputs=["output"],
        )
        input = np.array([[1, 2], [3, 4]], dtype=dtype)
        weights = np.array([[[1, 0], [0, 1]], [[0, 1], [1, 0]]], dtype=dtype)
        group_indices = np.array([[0, 1], [1, 0]], dtype=index_dtype)
        bias = np.array([[0.5, 1], [1.5, 2]], dtype=dtype)
        output = grouped_matmul(input, weights, group_indices, bias)

        expect(
            node,
            inputs=[input, weights, group_indices, bias],
            outputs=[output],
            name=f"test_groupedmatmul_{np.dtype(dtype).name}_{np.dtype(index_dtype).name}",
            opset_imports=[onnx.helper.make_opsetid("", 29)],
        )
```

#### default

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array(
    [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]], dtype=np.float32
)
weights = np.array(
    [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [0, 1]]],
    dtype=np.float32,
)
group_indices = np.array([[0], [1], [0], [1]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _without_bias_explicit_empty_input

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices", ""],
    outputs=["output"],
)
input = np.array(
    [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]], dtype=np.float32
)
weights = np.array(
    [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [0, 1]]],
    dtype=np.float32,
)
group_indices = np.array([[0], [1], [0], [1]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_without_bias_explicit_empty_input",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _with_bias

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices", "bias"],
    outputs=["output"],
)
input = np.array([[1, 0], [0, 1]], dtype=np.float32)
weights = np.array(
    [[[1, 0], [0, 1]], [[0, 1], [1, 0]], [[1, 1], [0, 0]]],
    dtype=np.float32,
)
group_indices = np.array([[0, 1], [2, 0]], dtype=np.int64)
bias = np.array([[0.1, 0.2], [0.3, 0.0], [0.5, 0.5]], dtype=np.float32)
output = grouped_matmul(input, weights, group_indices, bias)

expect(
    node,
    inputs=[input, weights, group_indices, bias],
    outputs=[output],
    name="test_groupedmatmul_with_bias",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _with_unused_group

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=np.float32)
weights = np.arange(12, dtype=np.float32).reshape(3, 2, 2)
group_indices = np.array([[0], [0], [2], [2]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_with_unused_group",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _single_group

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
weights = np.array([[[1, 2, 3], [4, 5, 6]]], dtype=np.float32)
group_indices = np.zeros((3, 1), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_single_group",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _zero_selections

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.arange(6, dtype=np.float32).reshape(2, 3)
weights = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
group_indices = np.empty((2, 0), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_zero_selections",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```

#### _zero_tokens

```python
import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.empty((0, 3), dtype=np.float32)
weights = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
group_indices = np.empty((0, 2), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_zero_tokens",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)
```