#!/usr/bin/env python3
"""Test quantized shard shapes."""

import torch
from pathlib import Path


def test_shapes():
    print("Loading quantized test shard...")
    ckpt = torch.load("/data/models/test_quantize/model.safetensors", map_location="cpu")
    
    print(f"Loaded {len(ckpt)} tensors\n")
    
    # Check for mismatched shapes
    mismatch_count = 0
    for key, tensor in list(ckpt.items())[:20]:  # Check first 20
        if isinstance(tensor, torch.Tensor):
            print(f"✓ {key} -> {tuple(tensor.shape)}")
        elif hasattr(tensor, 'shape'):
            print(f"✓ {key} -> {tensor.shape}")
        else:
            print(f"? {key} -> {type(tensor)}")
    
    print(f"\nTotal tensors: {len(ckpt)}")
    print("If all shapes look correct, the quantization is working!")


if __name__ == "__main__":
    test_shapes()