File size: 1,420 Bytes
dd7a16d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | import json, struct
import numpy as np
import mlx.core as mx
def _load(path, key):
with open(path, 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0]
hdr = json.loads(f.read(n))
e = hdr[key]
start, end = e['data_offsets']
f.seek(8 + n + start)
raw = f.read(end - start)
return e, raw
def get_tensor(path, key, dtype=None):
"""Bit-accurate tensor reader. BF16 is bit-reinterpreted (uint16 bits -> bf16)."""
e, raw = _load(path, key)
t = e['dtype']
if t == 'BF16':
u16 = np.frombuffer(raw, dtype=np.uint16)
f32 = np.left_shift(u16.astype(np.uint32), 16).view(np.float32).copy()
a = mx.array(f32)
if dtype in (None, 'bf16'):
# exact bf16 round-trip: bf16 -> f32 is lossless
return a.reshape(e['shape']) if dtype is None else a.astype(mx.bfloat16).reshape(e['shape'])
return a.reshape(e['shape'])
if t == 'F16':
return mx.array(np.frombuffer(raw, dtype=np.float16).copy()).reshape(e['shape'])
if t == 'F32':
return mx.array(np.frombuffer(raw, dtype=np.float32).copy()).reshape(e['shape'])
if t in ('U32', 'I32'):
return mx.array(np.frombuffer(raw, dtype=np.uint32).copy()).reshape(e['shape'])
raise ValueError(t)
def header(path):
with open(path, 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0]
return json.loads(f.read(n))
|