def from_float(cls, mod: _nn.Module) -> _nn.Module:
        r"""Create a palettization module from a float module or qparams_dict

        Args: `mod` a float module, either produced by torch.quantization utilities
        or directly from user
        """
        assert type(mod) == cls._FLOAT_MODULE, (
            "qat." + cls.__name__ + ".from_float only works for " + cls._FLOAT_MODULE.__name__
        )
        assert hasattr(mod, "qconfig"), "Input float module must have qconfig defined"
        assert mod.qconfig, "Input float module must have a valid qconfig"

        qconfig = mod.qconfig
        qat = cls(
            mod.embed_dim,
            mod.num_heads,
            mod.dropout,
            batch_first=mod.batch_first,
            bias=hasattr(mod, "in_proj_bias"),
            add_bias_kv=mod.bias_k is not None,
            add_zero_attn=mod.add_zero_attn,
            kdim=mod.kdim,
            vdim=mod.vdim,
        )
        qat.qconfig = qconfig
        if not qat._qkv_same_embed_dim:
            qat.q_proj_weight_fake_quant = qconfig.weight()
            qat.k_proj_weight_fake_quant = qconfig.weight()
            qat.v_proj_weight_fake_quant = qconfig.weight()
        else:
            qat.in_proj_weight_fake_quant = qconfig.weight()

        qat.out_proj.weight_fake_quant = qconfig.weight()

        for name, param in mod.named_parameters(recurse=False):
            setattr(qat, name, param)

        for name, param in mod.out_proj.named_parameters(recurse=False):
            setattr(qat.out_proj, name, param)

        return qat
Create a palettization module from a float module or qparams_dict

        Args: `mod` a float module, either produced by torch.quantization utilities
        or directly from user
        
coremltools/optimize/torch/palettization/_supported_modules.py