tempus_bench.models.toto.toto.model.util

`RMSNorm`

chronax.util.RMSNorm · inherits torch.nn.Module

Wraps xFormers' rms_norm for eval/frozen mode, and does a Python fallback for train mode.

`init(self, dim: int, include_weight: bool = True, eps: float = 1e-8)`

Parameter	Type	Default	Description
`dim`	`int`	-	(undocumented)
`include_weight`	`bool`	`True`	(undocumented)
`eps`	`float`	`1e-8`	(undocumented)

`forward(self, x: torch.Tensor)`

Parameter	Type	Default	Description
`x`	`torch.Tensor`	-	(undocumented)

`increment_and_forward_(self, x: torch.Tensor, y: torch.Tensor)`

If you need the fused addition with RMS norm, do the same check here.

Parameter	Type	Default	Description
`x`	`torch.Tensor`	-	(undocumented)
`y`	`torch.Tensor`	-	(undocumented)

`make_batched_block_mask`

chronax.util.make_batched_block_mask

Parameter	Type	Default	Description
`t`	`torch.Tensor`	-	(undocumented)

`KVCache`

chronax.util.KVCache

Key/Value cache for storing intermediate attention values during multistep inference. Only stores KV cache for timewise layers, skipping spacewise layers.

`init(self, batch_size: int, num_variates: int, transformer_layers: List["TransformerLayer"], num_layers: int, embed_dim: int, num_heads: int, max_seq_len: int, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, use_memory_efficient_attention: bool = True)`

Parameter	Type	Default	Description
`batch_size`	`int`	-	(undocumented)
`num_variates`	`int`	-	(undocumented)
`transformer_layers`	`List["TransformerLayer"]`	-	(undocumented)
`num_layers`	`int`	-	(undocumented)
`embed_dim`	`int`	-	(undocumented)
`num_heads`	`int`	-	(undocumented)
`max_seq_len`	`int`	-	(undocumented)
`device`	`torch.device`	`torch.device('cpu')`	(undocumented)
`dtype`	`torch.dtype`	`torch.float32`	(undocumented)
`use_memory_efficient_attention`	`bool`	`True`	(undocumented)

`getitem(self, layer_idx: int) -> KV`

Parameter	Type	Default	Description
`layer_idx`	`int`	-	(undocumented)

`current_len(self, cache_idx: int) -> int`

Parameter	Type	Default	Description
`cache_idx`	`int`	-	(undocumented)

`seq_len(self, layer_idx: int) -> int`

Parameter	Type	Default	Description
`layer_idx`	`int`	-	(undocumented)

`append(self, layer_idx: int, kv: KV)`

Parameter	Type	Default	Description
`layer_idx`	`int`	-	(undocumented)
`kv`	`KV`	-	(undocumented)

RMSNorm

__init__(self, dim: int, include_weight: bool = True, eps: float = 1e-8)

forward(self, x: torch.Tensor)

increment_and_forward_(self, x: torch.Tensor, y: torch.Tensor)

make_batched_block_mask

KVCache

__init__(self, batch_size: int, num_variates: int, transformer_layers: List["TransformerLayer"], num_layers: int, embed_dim: int, num_heads: int, max_seq_len: int, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, use_memory_efficient_attention: bool = True)

__getitem__(self, layer_idx: int) -> KV

current_len(self, cache_idx: int) -> int

seq_len(self, layer_idx: int) -> int

append(self, layer_idx: int, kv: KV)

reset(self)