Apple MLX weight files (npy / npz / safetensors) and mlx-lm quantization

def read_mlx(path, format=None, return_metadata=False):
    if format is None:
        format = path.name.rsplit(".", 1)[-1]          # file objects use the .name attribute
    if format == "npy":
        if return_metadata: raise Invalid("metadata not supported")
        return read_mlx_npy(path)
    if format == "npz":
        if return_metadata: raise Invalid("metadata not supported")
        z = ZipFile(path)
        out = {}
        for name in z.namelist():
            key = name[:-4] if name.endswith(".npy") else name
            out[key] = read_mlx_npy(z.open(name))
        return out
    if format == "safetensors":
        arrays, meta = read_safetensors(path)          # safetensors spec, plus MLX-003 (header length) and MLX-004 (data_offsets)
        return (arrays, meta) if return_metadata else arrays
    if format == "gguf":
        return read_mlx_gguf(path, return_metadata)    # this spec's GGUF reader: reverse ne, copy plain types, Q4_0/Q4_1/Q8_0 become affine triplets
    raise Invalid("unknown format")

def read_mlx_npy(f):
    magic, major, minor = read_exact(f, 8)
    if magic != b"\x93NUMPY" or major not in (1, 2):
        raise Invalid("bad npy header")
    hlen = u16(f) if major == 1 else u32(f)            # host endian, not swapped
    header = read_exact(f, hlen).decode("ascii")
    descr = header[11:14]                              # bytes [11, 14); assumes the 3-character descr this writer emits (MLX-010)
    dtype = DTYPE_FROM_DESCR[descr]                    # 'V2' -> bfloat16
    fortran = header[34] == "T"
    shape = parse_ints_between_last_parens(header)     # rejects negatives
    raw = read_exact(f, prod(shape) * itemsize(dtype))
    if descr[0] == ">" and host_is_little():
        raw = byteswap(raw, itemsize(dtype))
    if fortran:
        shape = list(reversed(shape))
        arr = Tensor(dtype, shape, raw)
        arr = transpose(arr)
    else:
        arr = Tensor(dtype, shape, raw)
    return arr

def read_mlx_gguf(path, return_metadata=False):
    weights, meta = parse_gguf(path)               # GGUF spec
    out = {}
    for name, t in weights.items():
        shape = list(reversed(t.ne))               # ne[0] was innermost
        if t.type in (F32, F16, I8, I16, I32):
            out[name] = Tensor(GGUF_TO_MLX[t.type], shape, t.data)
        elif t.type in (Q4_0, Q4_1, Q8_0):
            w, s, b = gguf_quant_to_affine(t)       # below
            prefix = name[:-7]                      # drop the last 7 characters, assumed to be ".weight" (MLX-012)
            out[name] = w; out[prefix+".scales"] = s; out[prefix+".biases"] = b
        elif t.type in (BF16, Q2_K, Q4_K, Q6_K):
            out[name] = Tensor(float16, shape, gguflib_to_f16(t))   # lossy
        else:
            raise Invalid("gguf_tensor_to_f16 failed")
    return (out, meta) if return_metadata else out

def gguf_quant_to_affine(t):
    # blocks of 32 along the innermost ggml axis, which is MLX's last axis
    if t.type == Q4_0:          # block: f16 scale, then 16 bytes of nibbles
        # ggml byte j = elem[j] | elem[j+16] << 4
        # MLX wants consecutive nibbles, low nibble = earlier element
        scale = f16_at(block, 0)
        codes = repack_q4_nibbles(block[2:18])      # 32 codes in 0..15, then pack_pow2(..., 4)
        bias = -8 * scale
        bits = 4
    elif t.type == Q4_1:        # f16 scale, f16 bias, 16 nibble bytes
        scale, bias = f16_at(block, 0), f16_at(block, 2)
        codes = repack_q4_nibbles(block[4:20]); bits = 4
    else:                       # Q8_0: f16 scale, 32 int8 codes
        scale = f16_at(block, 0); bias = -128 * scale
        codes = (int8_bytes ^ 0x80).astype(uint8)    # signed -128..127 -> 0..255
        bits = 8
    # dequant is then scale * code + bias, which restores
    # Q4_0: scale * (nibble - 8); Q4_1: scale * nibble + min; Q8_0: scale * int8
    return pack_pow2(codes, bits), scales_f16, biases_f16