def test_multiple_parameters_in_same_layer(
        self, compute_unit, n_bits, minimum_deployment_target, frontend
    ):
        """Test one layer has multiple parameters (such as weight and bias in a linear layer)"""
        if n_bits == 4 and minimum_deployment_target < ct.target.iOS18:
            pytest.skip("Sub-byte quantization is only supported since iOS18.")

        class Model(torch.nn.Module):
            def __init__(self):
                super(Model, self).__init__()
                self.linear_1 = torch.nn.Linear(16, 32)
                self.linear_2 = torch.nn.Linear(32, 64)

            def forward(self, x):
                return self.linear_2(self.linear_1(x))

        model = Model().eval()
        with torch.no_grad():
            fake_weight_scale = 2 if n_bits == 4 else 40
            model.linear_2.weight = torch.nn.Parameter(
                torch.from_numpy(
                    np.ones_like(model.linear_2.weight.detach().numpy()) * fake_weight_scale
                ).float()
            )
            model.linear_2.bias = torch.nn.Parameter(
                torch.from_numpy(
                    np.ones_like(model.linear_2.bias.detach().numpy()) * fake_weight_scale
                ).float()
            )

        # Register buffers for both weight and bias for linear_2 layer.
        weight_scale = np.array([2.0], dtype=np.float32).reshape(1, 1)
        bias_scale = np.array([2.0], dtype=np.float32)
        model.linear_2.register_buffer("_COREML_/weight/compression_type", torch.tensor([3]))
        model.linear_2.register_buffer("_COREML_/weight/quantization_n_bits", torch.tensor(n_bits))
        model.linear_2.register_buffer(
            "_COREML_/weight/quantization_scale", torch.from_numpy(weight_scale)
        )
        model.linear_2.register_buffer("_COREML_/bias/compression_type", torch.tensor([3]))
        model.linear_2.register_buffer("_COREML_/bias/quantization_n_bits", torch.tensor(n_bits))
        model.linear_2.register_buffer(
            "_COREML_/bias/quantization_scale", torch.from_numpy(bias_scale)
        )
        model.register_buffer("_COREML_/metadata_version", torch.tensor(2))

        res = self.run_compare_torch(
            [(8, 16)],
            model,
            minimum_deployment_target=minimum_deployment_target,
            compute_unit=compute_unit,
            frontend=frontend,
            converter=ct.convert,
        )
        main_func = res[1]._mil_program.functions["main"]

        quantize_op_type = (
            "constexpr_blockwise_shift_scale"
            if minimum_deployment_target >= ct.target.iOS18
            else "constexpr_affine_dequantize"
        )
        # Only the linear_2 layer got quantized based on registered buffers.
        linear_ops = main_func.find_ops(op_type="linear")
        assert linear_ops[0].weight.op.op_type == "const"
        assert linear_ops[0].bias.op.op_type == "const"
        if frontend == TorchFrontend.EXECUTORCH:
            # In EXECUTORCH, the second linear layer is represented by `matmul` and `add` op.
            matmul_op = main_func.find_ops(op_type="matmul")[0]
            add_op = main_func.find_ops(op_type="add")[0]
            assert matmul_op.y.op.op_type == quantize_op_type
            assert add_op.x.op.op_type == quantize_op_type
        else:
            assert linear_ops[1].weight.op.op_type == quantize_op_type
            assert linear_ops[1].bias.op.op_type == quantize_op_type

        quantize_ops = main_func.find_ops(op_type=quantize_op_type)
        assert len(quantize_ops) == 2
        for quantize_op in quantize_ops:
            if minimum_deployment_target >= ct.target.iOS18:
                assert types.builtin_to_string(quantize_op.data.dtype) == f"uint{n_bits}"
            else:
                assert types.builtin_to_string(quantize_op.quantized_data.dtype) == f"uint{n_bits}"
Test one layer has multiple parameters (such as weight and bias in a linear layer)
coremltools/converters/mil/frontend/torch/test/test_torch_quantization_ops.py