Skip to content

State Management

Utilities for flattening, unflattening, and detaching nested model state.

StateSpec

Bases: NamedTuple

Batch-agnostic specification for a nested model state structure.

Stores shape templates with -1 at the batch dimension, enabling flatten/unflatten at any batch size.

state_size property 

state_size: int

Total flat state dimension (sum of all leaf widths).

detach_state 

detach_state(state)

Recursively detach tensors from the computation graph.

Source code in tsfast/models/_core/state.py 
26
27
28
29
30
31
32
33
34
35
36
def detach_state(state):
    """Recursively detach tensors from the computation graph."""
    if state is None:
        return None
    if isinstance(state, torch.Tensor):
        return state.detach()
    if isinstance(state, (list, tuple)):
        return type(state)(detach_state(s) for s in state)
    if isinstance(state, dict):
        return {k: detach_state(v) for k, v in state.items()}
    return state

discover_state_spec 

discover_state_spec(model: Module, n_in: int, device: str | device = 'cpu') -> StateSpec

Probe a stateful model to discover its state structure.

Forwards two batches of different sizes and compares the resulting state shapes to identify the batch dimension. Templates store -1 at the batch dim so the spec is batch-agnostic.

Parameters:

Name Type Description Default
model Module

stateful model returning (output, state)

 required
n_in int

number of input features

 required
device str | device

device for probe tensors

 'cpu'
Source code in tsfast/models/_core/state.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def discover_state_spec(model: nn.Module, n_in: int, device: str | torch.device = "cpu") -> StateSpec:
    """Probe a stateful model to discover its state structure.

    Forwards two batches of different sizes and compares the resulting state
    shapes to identify the batch dimension.  Templates store ``-1`` at the
    batch dim so the spec is batch-agnostic.

    Args:
        model: stateful model returning ``(output, state)``
        n_in: number of input features
        device: device for probe tensors
    """
    was_training = model.training
    model.eval()
    try:
        with torch.no_grad():
            _, state_a = model(torch.zeros(2, 1, n_in, device=device), state=None)
            _, state_b = model(torch.zeros(3, 1, n_in, device=device), state=None)
    finally:
        if was_training:
            model.train()

    leaves_a, tree_spec = pytree.tree_flatten(state_a)
    leaves_b, _ = pytree.tree_flatten(state_b)

    templates = []
    widths = []
    for la, lb in zip(leaves_a, leaves_b):
        sa, sb = la.shape, lb.shape
        assert len(sa) == len(sb), f"Rank mismatch: {sa} vs {sb}"
        template = []
        width = 1
        batch_found = False
        for da, db in zip(sa, sb):
            if da != db and not batch_found:
                template.append(-1)
                batch_found = True
            else:
                template.append(da)
                width *= da
        assert batch_found, f"No batch dim found: shapes {sa} vs {sb} are identical"
        templates.append(tuple(template))
        widths.append(width)

    return StateSpec(tree_spec, tuple(templates), tuple(widths))

build_spec_from_state 

build_spec_from_state(state, batch_size: int) -> StateSpec

Build a StateSpec from a single state sample by inferring the batch dimension.

The batch dimension is identified as the first dimension matching batch_size. Ambiguous when a non-batch dimension equals batch_size.

Source code in tsfast/models/_core/state.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def build_spec_from_state(state, batch_size: int) -> StateSpec:
    """Build a StateSpec from a single state sample by inferring the batch dimension.

    The batch dimension is identified as the first dimension matching ``batch_size``.
    Ambiguous when a non-batch dimension equals ``batch_size``.
    """
    leaves, tree_spec = pytree.tree_flatten(state)
    templates = []
    widths = []
    for leaf in leaves:
        template = []
        width = 1
        batch_found = False
        for d in leaf.shape:
            if d == batch_size and not batch_found:
                template.append(-1)
                batch_found = True
            else:
                template.append(d)
                width *= d
        assert batch_found, f"No dim matching batch_size={batch_size} in shape {leaf.shape}"
        templates.append(tuple(template))
        widths.append(width)
    return StateSpec(tree_spec, tuple(templates), tuple(widths))

flatten_state 

flatten_state(state, batch_size: int) -> Tensor

Flatten arbitrary nested state to a single [B, D] tensor.

Source code in tsfast/models/_core/state.py 
112
113
114
115
116
def flatten_state(state, batch_size: int) -> Tensor:
    """Flatten arbitrary nested state to a single ``[B, D]`` tensor."""
    leaves, _ = pytree.tree_flatten(state)
    flat_leaves = [leaf.reshape(batch_size, -1) for leaf in leaves]
    return torch.cat(flat_leaves, dim=-1)

unflatten_state 

unflatten_state(flat: Tensor, spec: StateSpec)

Reconstruct nested state from a flat [B, D] tensor + StateSpec.

Source code in tsfast/models/_core/state.py 
119
120
121
122
123
124
125
126
127
def unflatten_state(flat: Tensor, spec: StateSpec):
    """Reconstruct nested state from a flat ``[B, D]`` tensor + StateSpec."""
    B = flat.shape[0]
    splits = flat.split(spec.widths, dim=-1)
    leaves = []
    for s, template in zip(splits, spec.templates):
        shape = tuple(B if d == -1 else d for d in template)
        leaves.append(s.reshape(shape).contiguous())
    return pytree.tree_unflatten(leaves, spec.tree_spec)