Python module
rotary_embedding
The rope embedding used within the model.
DeepseekYarnRopeScalingParams
class max.nn.rotary_embedding.DeepseekYarnRopeScalingParams(scaling_factor: float, original_max_position_embeddings: int, beta_fast: int, beta_slow: int, mscale: float, mscale_all_dim: float)
-
Parameters:
beta_fast
beta_fast*: int*
Fast interpolation rate.
beta_slow
beta_slow*: int*
Slow interpolation rate.
mscale
mscale*: float*
Scaling factor for middle frequencies.
mscale_all_dim
mscale_all_dim*: float*
Scaling factor applied to all dimensions.
original_max_position_embeddings
original_max_position_embeddings*: int*
Original maximum sequence length during training.
scaling_factor
scaling_factor*: float*
Scaling factor for frequency interpolation.
DeepseekYarnRotaryEmbedding
class max.nn.rotary_embedding.DeepseekYarnRotaryEmbedding(dim, n_heads, theta, max_seq_len, device, head_dim=None, _freqs_cis=None, interleaved=True, scaling_params=None)
Deepseek’s YaRN (Yet another RoPE eNhancement) Rotary Position Embedding layer.
Unlike Llama3RotaryEmbedding, the dim argument here is the rope dimension of the model, not the hidden dimension.
-
Parameters:
compute_scale()
compute_scale(user_scale=None)
freqs_cis_base()
freqs_cis_base()
Computes the frequency tensor for complex exponentials (cis) for a given seq_len. Tensor is scaled with theta parameter. Required to apply Rotary Position Embedding (RoPE) to tensor. See ‘Roformer: Enhanced Transformer with Rotary Embedding’ (arxiv.org/pdf/2104.09864).
-
Returns:
-
The frequency tensor for complex exponentials with shape (max_seq_len, rope_dim // 2, 2)
-
Return type:
scaling_params
scaling_params*: DeepseekYarnRopeScalingParams | None* = None
LinearScalingParams
class max.nn.rotary_embedding.LinearScalingParams(factor: float)
-
Parameters:
-
factor (
float)
factor
factor*: float*
Main scaling factor for the frequency components of the rope.
Llama3RopeScalingParams
class max.nn.rotary_embedding.Llama3RopeScalingParams(factor: float, low_freq_factor: float, high_freq_factor: float, orig_max_position: int)
-
Parameters:
factor
factor*: float*
Main scaling factor for the frequency components of the rope.
high_freq_factor
high_freq_factor*: float*
Factor to scale the high frequency components of the rope.
low_freq_factor
low_freq_factor*: float*
Factor to scale the low frequency components of the rope.
orig_max_position
orig_max_position*: int*
The original maximum position length supported by the model.
Llama3RotaryEmbedding
class max.nn.rotary_embedding.Llama3RotaryEmbedding(dim, n_heads, theta, max_seq_len, device, head_dim=None, _freqs_cis=None, interleaved=True, scaling_params=None)
RotaryEmbedding for Llama3 that takes rope scaling into account.
-
Parameters:
scaling_params
scaling_params*: Llama3RopeScalingParams | None* = None
Scaling parameters to enable llama to function with a longer context length.
OptimizedRotaryEmbedding
class max.nn.rotary_embedding.OptimizedRotaryEmbedding(dim, n_heads, theta, max_seq_len, device, head_dim=None, _freqs_cis=None, interleaved=True)
Optimized version of RotaryEmbedding using 2D frequency tensor representation.
-
Parameters:
freqs_cis
property freqs_cis
RotaryEmbedding
class max.nn.rotary_embedding.RotaryEmbedding(dim, n_heads, theta, max_seq_len, device, head_dim=None, _freqs_cis=None, interleaved=True)
RotaryEmbedding layer to calculate and apply the frequency tensor for complex exponentials.
-
Parameters:
compute_scale()
compute_scale(user_scale=None)
device
device*: DeviceRef*
dim
dim*: int*
freqs_cis
property freqs_cis*: TensorValue*
freqs_cis_base()
freqs_cis_base()
Computes the frequency tensor for complex exponentials (cis) for a given seq_len. Tensor is scaled with theta parameter. Required to apply Rotary Position Embedding (RoPE) to tensor. See ‘Roformer: Enhanced Transformer with Rotary Embedding’ (arxiv.org/pdf/2104.09864).
-
Returns:
-
The frequency tensor for complex exponentials with shape (max_seq_len * 2, head_dim / 2, 2)
-
Return type:
head_dim
head_dim = dim // n_heads if not specified in the config.
interleaved
interleaved*: bool* = True
max_seq_len
max_seq_len*: int*
The maximum sequence length for model’s input.
n_heads
n_heads*: int*
theta
theta*: float*
Hyperparameter used to control the frequency scaling of the sinusoidal components of the embeddings.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!