"""Print tensor and segment metadata from an H3 direct-runner capture."""

import argparse

import torch


parser = argparse.ArgumentParser()
parser.add_argument("capture_dir")
args = parser.parse_args()

for name in ("input.pt", "output.pt"):
    payload = torch.load(f"{args.capture_dir}/{name}", map_location="cpu", weights_only=False)
    print(name)
    for key, value in payload.items():
        if isinstance(value, torch.Tensor):
            print(f"  {key}: shape={tuple(value.shape)} dtype={value.dtype} finite={torch.isfinite(value).all().item()}")
        else:
            print(f"  {key}: {value}")