MemoryPlan

MemoryPlan()

Top-level memory plan for tensors.

Provisional: importable, but outside the compatibility promise; a minor release can change or remove it.

Every tensor allocation is a slot in some arena. Three arena maps partition storage by role:

:attr:arena_usages — writable scratch arenas (one per writable memory bank).
:attr:persistent_arenas — writable persistent arenas (one per writable memory bank that received a persistent tensor). Separate from scratch so the scratch arena stays purely transient and may be aliased across models.
:attr:constant_arenas — constant arenas. Two shapes share the same map:
- Cold: arena.source_memory == arena.memory. The arena buffer itself is the read-only blob in cold storage; kernels read it directly. No hydration required.
- Staged: arena.source_memory != arena.memory. The arena buffer is a writable runtime copy; a separate source blob in arena.source_memory is hydrated into it before the first model_run.

tensor_allocs: dict[str, TensorAllocation] = Field(..., default_factory=dict, description='Tensor allocations')

Tensor allocations. Each carries a :class:TensorBinding recording (role, memory, offset, source_memory).

arena_usages: dict[MemoryType, ArenaUsage] = Field(..., default_factory=dict, description='Memory arena usages')

Scratch arenas, keyed by MemoryType.

constant_arenas: dict[MemoryType, ArenaUsage] = Field(..., default_factory=dict, description='Per-bank constant arenas (one per runtime memory).')

Constant arenas, keyed by runtime (destination) MemoryType. Per-arena single-source-memory invariant holds (every constant in the same arena has the same source_memory) so cold arenas and staged arenas are both guaranteed to be a single contiguous source blob.

persistent_arenas: dict[MemoryType, ArenaUsage] = Field(..., default_factory=dict, description='Per-bank persistent (resource-variable) arenas.')

Persistent arenas, keyed by writable MemoryType.

tensor_lifetimes: dict[TensorId, TensorLifetime] = Field(..., default_factory=dict, description='Tensor lifetimes')

Tensor lifetimes.

get_allocation(tensor_id: str) -> TensorAllocation

Get the allocation metadata for a given tensor ID.

Parameters

Parameters of get_allocation
Name	Type	Default	Description
`tensor_id`	`str`	Required	The tensor ID to look up.

Returns

Returns of get_allocation
Value	Type	Description
`TensorAllocation`	`TensorAllocation`	The allocation metadata for the tensor.

Raises

Errors raised by get_allocation
Type	Description
`KeyError`	If the tensor ID is not found in the allocations.