Core

Core classes provide the state-management and parameter-management machinery used by every traceTorch layer.

Model

class tracetorch.core.Model(*args: Any, **kwargs: Any)

Bases: Module

Manage hidden states and parameter caches across a model hierarchy.

Weight checkpoints use native PyTorch state_dict(). Hidden-state checkpoints use save_states()/load_states() and contain no parameters.

save_states() → Dict[str, torch.Tensor]

Return detached copies of initialized states, keyed by layer path.

None states are omitted. The result can be saved with torch.save(); neither raw weights nor compiled parameter buffers are included.

load_states(states: Dict[str, torch.Tensor], strict: bool = True, device=None) → None

Load a hidden-state checkpoint without modifying model parameters.

Strict mode rejects missing and extra state names. Existing shapes and declared trailing shapes are checked. Loaded tensors are detached and copied, so subsequent updates do not mutate the checkpoint dictionary. Device defaults to the layer’s parameters, buffers, or saved tensor.

reset_states() → None

Set all hidden states to None for lazy initialization next forward.

detach_states() → None

Detach every hidden state while preserving its numerical value.

compile_parameters() → None

Cache activated traceTorch parameters across the model for inference.

decompile_parameters() → None

Delete caches and restore intended parameter learnability.

Layer

class tracetorch.core.Layer(*args: Any, **kwargs: Any)

Bases: Module

Base class for stateful layers and optionally transformed parameters.

States are ordinary tensor attributes, separate from PyTorch weight checkpoints. Each state declares its own trailing shape and input slice. Parameters retain their raw nn.Parameter throughout compilation.

define_parameter(name: str, value: torch.Tensor, learnable: bool = True, initialization_fn: Callable | None = None, activation_fn: Callable | None = None) → None

Define a raw parameter and its optional runtime transformation.

value is a tensor of any shape. initialization_fn converts it into the raw representation once (e.g. logit for a sigmoid decay). self.raw_<name> is always an nn.Parameter; self.<name> returns its activated value, or a detached cache while compiled. Only the raw parameter is included in state_dict().

set_parameter_learnable(name: str, learnable: bool) → None

Set intended learnability, including while temporarily compiled.

Use this instead of changing raw_<name>.requires_grad directly: compilation freezes the raw parameter without changing this intention.

define_state(name: str, shape: Tuple[int, ...], dim: int | None = -1) → None

Declare a state, initially None, with its own allocation rule.

Allocation uses reference.shape[:dim] + shape. Negative dim drops trailing reference dimensions; None drops none and appends shape to the full reference shape. Zero and positive indices are rejected. An empty shape is allowed.

detach_state(state_name: str) → None

Detach an existing state from its computation graph.

detach_states() → None

Detach all states.

reset_state(state_name: str) → None

Set a declared state to None for lazy reinitialization.

reset_states() → None

Reset all states to None.

zero_state(state_name: str, reference_tensor: torch.Tensor) → None

Allocate zeros from a state’s shape rule, only if it is None.

The reference supplies the leading instance dimensions, dtype and device. Existing states are not overwritten or silently resized. Custom layers may override this method for nonzero initialization.

zero_states(reference_tensor: torch.Tensor) → None

Lazily allocate all states using the same reference tensor.

compile_parameter(name: str) → None

Cache an activated parameter for inference and freeze its raw value.

The cache is a nonpersistent buffer. Raw parameter identity and intended learnability are preserved; repeated compilation is a no-op.

compile_parameters() → None

Compile every parameter defined through define_parameter.

decompile_parameter(name: str) → None

Delete a parameter cache and restore intended learnability.

decompile_parameters() → None

Decompile every parameter defined through define_parameter.