50 lines
1.9 KiB
Python
50 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract an upstream Kimodo NPZ capture into checked raw F32 tensors.
|
|
|
|
The raw files are intentionally dependency-free inputs for the C++ parity
|
|
tests and demo tooling. ``shapes.json`` preserves every original array shape
|
|
and dtype; integral arrays stay integral while floating arrays are canonical
|
|
little-endian F32.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("input", type=Path, help="upstream .npz capture")
|
|
parser.add_argument("output", type=Path, help="output fixture directory")
|
|
args = parser.parse_args()
|
|
|
|
if args.input.suffix != ".npz" or not args.input.is_file():
|
|
raise SystemExit("input must be an existing .npz file")
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
shapes: dict[str, dict[str, object]] = {}
|
|
with np.load(args.input, allow_pickle=False) as archive:
|
|
for name in sorted(archive.files):
|
|
value = np.ascontiguousarray(archive[name])
|
|
if value.dtype.kind == "f":
|
|
value = np.asarray(value, dtype="<f4")
|
|
suffix = ".f32"
|
|
elif value.dtype.kind in "iu":
|
|
value = np.asarray(value, dtype="<i8")
|
|
suffix = ".i64"
|
|
elif value.dtype.kind == "b":
|
|
value = np.asarray(value, dtype="u1")
|
|
suffix = ".u8"
|
|
else:
|
|
raise SystemExit(f"unsupported capture tensor {name}: {value.dtype}")
|
|
(args.output / f"{name}{suffix}").write_bytes(value.tobytes())
|
|
shapes[name] = {"shape": list(value.shape), "dtype": str(value.dtype), "file": f"{name}{suffix}"}
|
|
(args.output / "shapes.json").write_text(json.dumps(shapes, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|