diff --git a/backends/apple/metal/tests/test_modules.py b/backends/apple/metal/tests/test_modules.py index d1956fd2913..b841d0a47ab 100644 --- a/backends/apple/metal/tests/test_modules.py +++ b/backends/apple/metal/tests/test_modules.py @@ -250,6 +250,30 @@ def forward(self, x: torch.Tensor): } +# ------------------------------------------------------------------------- +class LinearWithBiasInt4(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(128, 256, bias=True) + + def forward(self, x: torch.Tensor): + return self.linear(x) + + +MODULE_REGISTRY["linear_bias_int4"] = { + "model_class": LinearWithBiasInt4, + "input_shapes": [(127, 128)], + "description": "Linear layer with bias and int4 quantization", + "qlinear": "fpa4w", + "qlinear_group_size": 32, + "compare_to_unquantized": False, + "atol_float32": 5e-2, + "rtol_float32": 5e-2, + "atol_bfloat16": 1e-1, + "rtol_bfloat16": 1e-1, + "skip": not TORCHAO_AVAILABLE, +} + # ------------------------------------------------------------------------- # Convolution Modules # -------------------------------------------------------------------------