TimeAwareRotaryEmbedding

chronax.rope.TimeAwareRotaryEmbedding · inherits RotaryEmbedding

A variant of the rotary position embedding that (optionally) uses the time index to compute the sinusoidal and cosine embeddings. This is useful for time series data, where the time index is the most important positional information.

__init__(self, *args, **kwargs)

Initializes the TimeAwareRotaryEmbedding, passing arguments to the base class. It also ensures that the freqs attribute, if present and registered as a parameter, is converted to a non-persistent buffer for compatibility with sharding mechanisms like FSDP.

Parameter	Type	Default	Description
*args	-	-	Positional arguments passed to RotaryEmbedding.__init__.
**kwargs	-	-	Keyword arguments passed to RotaryEmbedding.__init__.

rotate_queries_and_keys(self, q, k, seq_dim=None, seq_pos=None, seq_pos_offset=0)

This method is the same as the one on the base class, except it allows you to override the sequence position tensor with a custom one. It also removes the ability to cache the position encodings, since we have to compute them dynamically based on the timesteps in the input data.

Parameter	Type	Default	Description
q	torch.Tensor	-	(undocumented)
k	torch.Tensor	-	(undocumented)
seq_dim	Optional[int]	None	(undocumented)
seq_pos	Optional[Int[torch.Tensor, "... seq_len]]	None	Custom sequence position tensor to use instead of the default generated sequence.
seq_pos_offset	int	0	Offset applied to the sequence position tensor.

Returns: A tuple containing the rotated query and key tensors. ( torch.Tensor, torch.Tensor )

get_scale(self, t)

Adapted from the base class, but it knows how to handle when t has more than 1 dim (as is the case when we're using time-aware RoPE, and have a different sequence position vector for each time series).

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

Returns: torch.Tensor (The computed scale tensor).