File size: 1,479 Bytes
9aa90e0 | 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 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import itertools
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Mapping, MutableMapping
import torch
def substate(state: Mapping[str, torch.Tensor], key: str) -> dict[str, torch.Tensor]:
prefix = f"{key}."
prefix_len = len(prefix)
return {k[prefix_len:]: v for k, v in state.items() if k.startswith(prefix)}
def has_substate(state: Mapping[str, torch.Tensor], key: str) -> bool:
prefix = f"{key}."
return any(k.startswith(prefix) for k in state)
def pop_substate(state: MutableMapping[str, torch.Tensor], key: str) -> dict[str, torch.Tensor]:
prefix = f"{key}."
return {k.removeprefix(prefix): state.pop(k) for k in list(state) if k.startswith(prefix)}
def rename_substate(state: MutableMapping[str, torch.Tensor], key_from: str, key_to: str) -> None:
src_prefix = f"{key_from}." if key_from else ""
dst_prefix = f"{key_to}." if key_to else ""
for k in list(state):
if k.startswith(src_prefix):
state[dst_prefix + k.removeprefix(src_prefix)] = state.pop(k)
def indexed_substates(state: Mapping[str, torch.Tensor], key: str) -> list[dict[str, torch.Tensor]]:
result = []
for i in itertools.count():
s = substate(state, f"{key}.{i}")
if not s:
return result
result.append(s)
return []
|