TotoOutput
chronax.backbone.TotoOutput · inherits NamedTuple
Output of the Toto model. Contains the output distribution, the location parameters, and the scale parameters.
| Attribute | Type | Description |
|---|---|---|
distribution |
torch.distributions.Distribution |
(undocumented) |
loc |
jnp.ndarray |
(undocumented) |
scale |
jnp.ndarray |
(undocumented) |
TotoBackbone
chronax.backbone.TotoBackbone · inherits BaseForecaster
Toto (Timeseries-Optimized Transformer for Observability) is a transformer-based model for multivariate time series forecasting. It applies a patch embedding to the input data, followed by a transformer that alternates between time-wise and space-wise attention. The transformer is followed by a linear projection that maps the transformer output to the output distribution.
The output distribution can be a single distribution (e.g. Gaussian) or a mixture of distributions. If a mixture of distributions is used, the model will learn to predict the mixture weights as well as the parameters of the individual distributions.
__init__(self, patch_size, stride, embed_dim, num_layers, num_heads, mlp_hidden_dim, dropout, spacewise_every_n_layers, scaler_cls, output_distribution_classes, spacewise_first=True, output_distribution_kwargs=None, use_memory_efficient_attention=True, stabilize_with_global=True, scale_factor_exponent=10.0)
| Parameter | Type | Default | Description |
|---|---|---|---|
patch_size |
int |
- | Size of the patch to use for the patch embedding. |
stride |
int |
- | Stride to use for the patch embedding. |
embed_dim |
int |
- | Dimension of the model's latent space. |
num_layers |
int |
- | Number of transformer layers to use. |
num_heads |
int |
- | Number of attention heads to use in each self-attention layer. |
mlp_hidden_dim |
int |
- | Dimension of the hidden layer in the feedforward network. |
dropout |
float |
- | Dropout rate to use in the model. |
spacewise_every_n_layers |
int |
- | How many time-wise transformer layers to apply between each space-wise transformer layer. |
scaler_cls |
str |
- | Class to use for scaling the input data. |
output_distribution_classes |
list[str] |
- | List of classes to use for the output distribution. If a single class is provided, the model will output a single distribution. If multiple classes are provided, the model will output a learned mixture of distributions. |
spacewise_first |
bool |
True |
Whether to apply space-wise attention before time-wise attention. |
output_distribution_kwargs |
dict \| None |
None |
Keyword arguments to pass to the output distribution class. Note: this currently only works with a single output distribution class. |
use_memory_efficient_attention |
bool |
True |
Whether to use memory-efficient attention. If True, the model will use the memory-efficient from xFormers. |
stabilize_with_global |
bool |
True |
Whether to use global statistics to stabilize causal statistics by clamping extreme values. Only applies to causal scalers. |
scale_factor_exponent |
float |
10.0 |
Exponent that controls the allowed range of deviation from global scale for causal scalers. |
allocate_kv_cache(self, batch_size, num_variates, max_time_steps, device, dtype) -> KVCache
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
batch_size |
int |
- | (undocumented) |
num_variates |
int |
- | (undocumented) |
max_time_steps |
int |
- | (undocumented) |
device |
torch.device |
- | (undocumented) |
dtype |
torch.dtype |
- | (undocumented) |
Returns: KVCache
backbone(self, inputs, input_padding_mask, id_mask, kv_cache=None, scaling_prefix_length=None) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
inputs |
jnp.ndarray |
- | (undocumented) |
input_padding_mask |
jnp.ndarray |
- | (undocumented) |
id_mask |
jnp.ndarray |
- | (undocumented) |
kv_cache |
Optional[KVCache] |
None |
(undocumented) |
scaling_prefix_length |
Optional[int] |
None |
(undocumented) |
Returns: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]
forward(self, inputs, input_padding_mask, id_mask, kv_cache=None, scaling_prefix_length=None) -> TotoOutput
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
inputs |
jnp.ndarray |
- | (undocumented) |
input_padding_mask |
jnp.ndarray |
- | (undocumented) |
id_mask |
jnp.ndarray |
- | (undocumented) |
kv_cache |
Optional[KVCache] |
None |
(undocumented) |
scaling_prefix_length |
Optional[int] |
None |
(undocumented) |
Returns: TotoOutput
device
(undocumented)