| 48 | """load *.np file and return memmap and related tensor meta""" |
| 49 | |
| 50 | def parse_npy_header(fin: BinaryIO) -> dict[str, Any]: |
| 51 | start = fin.tell() |
| 52 | major, minor = np.lib.format.read_magic(fin) |
| 53 | if major == 1 and minor == 0: |
| 54 | read_header_fn = np.lib.format.read_array_header_1_0 |
| 55 | elif major == 2 and minor == 0: |
| 56 | read_header_fn = np.lib.format.read_array_header_2_0 |
| 57 | else: |
| 58 | raise ValueError( |
| 59 | f"unknown version {major}.{minor} when parsing npy header from {fn}" |
| 60 | ) |
| 61 | shape, is_fortran, dtype = read_header_fn(fin) |
| 62 | return { |
| 63 | "shape": shape, |
| 64 | "is_fortran": is_fortran, |
| 65 | "dtype": dtype, |
| 66 | "header_length": fin.tell() - start, |
| 67 | } |
| 68 | |
| 69 | meta_fn = fn + ".meta" |
| 70 | with open(meta_fn, "rb") as fin: |