File size: 1,747 Bytes
2ea2835
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
"""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",
    }