Function _write_safetensors

tests/test_oq.py:1242–1273

Write a minimal safetensors file from {name: np.ndarray} dict. Values can be np.ndarray (auto-dtype) or (raw_bytes, shape, sf_dtype) tuples for dtypes numpy doesn't support (F8_E4M3, F8_E8M0, I8).

(path, tensors)

def _write_safetensors(path, tensors):
    """Write a minimal safetensors file from {name: np.ndarray} dict.

    Values can be np.ndarray (auto-dtype) or (raw_bytes, shape, sf_dtype) tuples
    for dtypes numpy doesn't support (F8_E4M3, F8_E8M0, I8)."""
    import json
    import struct

    header = {}
    data_parts = []
    offset = 0
    dtype_map = {np.float16: "F16", np.float32: "F32", np.dtype("<f2"): "F16"}
    for name, val in tensors.items():
        if isinstance(val, tuple):
            raw, shape, sf_dtype = val
        else:
            raw = val.tobytes()
            shape = list(val.shape)
            sf_dtype = dtype_map.get(val.dtype, "F16")
        header[name] = {
            "dtype": sf_dtype,
            "shape": list(shape),
            "data_offsets": [offset, offset + len(raw)],
        }
        data_parts.append(raw)
        offset += len(raw)
    hdr_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(hdr_bytes)))
        f.write(hdr_bytes)
        for part in data_parts:
            f.write(part)


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")