`.typings/mlx_lm/models/bitlinear_layers.pyi`

27 lines
"""
This type stub file was generated by pyright.
"""

import mlx.nn as nn

def bitnet_quantize(model, quantization_config: dict): ...
def make_bitlinear_kernel():
    """
    Custom Metal kernel that performs matrix multiplication directly on
    packed weights and scales the output. This eliminates the need to
    store unpacked weights in memory.
    """

_bitlinear_kernel = ...

class BitLinear(nn.Module):
    """
    BitLinear module with memory-efficient weight handling.
    """
    def __init__(
        self, in_features, out_features, bias=..., invert_weight_scales=...
    ) -> None: ...
    def execute_matmul_kernel(self, x, packed_weights): ...
    def __call__(self, x):  # -> array:
        ...