| """Deterministic runtime identities for frozen representation modules.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| def module_state_sha256(module: nn.Module) -> str: |
| """Hash sorted tensor keys, metadata, and loaded tensor bytes.""" |
| digest = hashlib.sha256() |
| state = module.state_dict() |
| for key in sorted(state): |
| value = state[key] |
| if not isinstance(value, torch.Tensor): |
| raise TypeError( |
| f"state_dict entry {key!r} is not a tensor: " |
| f"{type(value).__name__}") |
| metadata = json.dumps( |
| { |
| "key": str(key), |
| "dtype": str(value.dtype), |
| "shape": [int(size) for size in value.shape], |
| }, |
| sort_keys=True, |
| separators=(",", ":"), |
| ).encode("utf-8") |
| digest.update(len(metadata).to_bytes(8, "little")) |
| digest.update(metadata) |
| cpu_bytes = ( |
| value.detach() |
| .to(device="cpu") |
| .contiguous() |
| .reshape(-1) |
| .view(torch.uint8) |
| .numpy() |
| ) |
| digest.update(int(cpu_bytes.nbytes).to_bytes(8, "little")) |
| digest.update(memoryview(cpu_bytes)) |
| return digest.hexdigest() |
|
|
|
|
| def representation_runtime_identity(module: nn.Module) -> dict: |
| """Return the strict identity used by frozen-encoder metadata checks.""" |
| return { |
| "wrapper_class": ( |
| f"{module.__class__.__module__}.{module.__class__.__qualname__}"), |
| "torch_version": str(torch.__version__), |
| "state_sha256": module_state_sha256(module), |
| "backend": "unknown", |
| "backend_version": "unknown", |
| } |
|
|
