AttentionAxis
chronax.attention.AttentionAxis
An enumeration defining the axis over which attention is computed.
| Name | Value |
|---|---|
TIME |
1 |
SPACE |
2 |
BaseMultiheadAttention
chronax.attention.BaseMultiheadAttention · inherits torch.nn.Module
(undocumented)
__init__(self, embed_dim: int, num_heads: int, dropout: float, rotary_emb: Optional[TimeAwareRotaryEmbedding], use_memory_efficient_attention: bool)
(undocumented)
| Parameter | Type | Default | Description |
|---|---|---|---|
embed_dim |
int |
- | (undocumented) |
num_heads |
int |
- | (undocumented) |
dropout |
float |
- | (undocumented) |
rotary_emb |
Optional[TimeAwareRotaryEmbedding] |
- | (undocumented) |
use_memory_efficient_attention |
bool |
- | (undocumented) |
forward(self, layer_idx: int, inputs: Float[torch.Tensor, "batch variate seq_len embed_dim"], attention_mask: Optional[Union[Bool[torch.Tensor, "batch_X_variate n_heads seq_len seq_len"], Bool[torch.Tensor, "batch_X_seq_len n_heads variate variate"]]] = None, kv_cache: Optional["KVCache"] = None) -> Float[torch.Tensor, "batch variate seq_len embed_dim"]
(undocumented)
| Parameter | Type | Default | Description |
|---|---|---|---|
layer_idx |
int |
- | (undocumented) |
inputs |
Float[torch.Tensor, "batch variate seq_len embed_dim"] |
- | (undocumented) |
attention_mask |
Optional[Union[Bool[torch.Tensor, "batch_X_variate n_heads seq_len seq_len"], Bool[torch.Tensor, "batch_X_seq_len n_heads variate variate"]]] |
None |
(undocumented) |
kv_cache |
Optional["KVCache"] |
None |
(undocumented) |
TimeWiseMultiheadAttention
chronax.attention.TimeWiseMultiheadAttention · inherits BaseMultiheadAttention
Computes standard multihead causal attention over the time axis. It does this by flattening out the variates along the batch dimension. It also applies rotary position embeddings to the query and key matrices in order to incorporate relative positional information.
SpaceWiseMultiheadAttention
chronax.attention.SpaceWiseMultiheadAttention · inherits BaseMultiheadAttention
Computes bidirectional multihead attention over the space axis (i.e. across variates within a multi-variate time series). This is done by flattening out the time axis along the batch dimension. This allows the model to attend to different variates at the same time point. By alternating between time-wise and space-wise attention, the model can learn both temporal and cross-variate dependencies in the data.
Unlike with time-wise attention, don't apply rotary embeddings here because we want cross-variate attention to be invariant to the order of the variates.
MultiHeadAttention
chronax.attention.MultiHeadAttention
Type alias for TimeWiseMultiheadAttention | SpaceWiseMultiheadAttention.