easydel.layers.rotary_embedding

class easydel.layers.rotary_embedding.DeepseekScalingRotaryEmbedding(*args: Any, **kwargs: Any)

Bases: Module

RotaryEmbedding extended with YaRN method.

class easydel.layers.rotary_embedding.DynamicNTKScalingRotaryEmbedding(*args: Any, **kwargs: Any)

Bases: RotaryEmbedding

RotaryEmbedding extended with Dynamic NTK scaling.

class easydel.layers.rotary_embedding.LinearScalingRotaryEmbedding(*args: Any, **kwargs: Any)

Bases: RotaryEmbedding

class easydel.layers.rotary_embedding.Llama3RotaryEmbedding(*args: Any, **kwargs: Any)

Bases: RotaryEmbedding

class easydel.layers.rotary_embedding.Phi3LongRoPEScaledRotaryEmbedding(*args: Any, **kwargs: Any)

Bases: Module

class easydel.layers.rotary_embedding.RopeConfig(rope_type: str = 'default', factor: Optional[float] = None, low_freq_factor: Optional[float] = None, high_freq_factor: Optional[float] = None, original_max_position_embeddings: Optional[int] = None, long_factor: Optional[float] = None, short_factor: Optional[float] = None, long_mscale: Optional[float] = None, short_mscale: Optional[float] = None)

Bases: Mapping

Configuration class for RoPE (Rotary Position Embedding) parameters.

factor: Optional[float] = None

classmethod from_dict(config_dict: Dict[str, Any]) → RopeConfig

Create a RopeConfig instance from a dictionary.

from_tuple()

high_freq_factor: Optional[float] = None

items() → a set-like object providing a view on D's items

keys() → a set-like object providing a view on D's keys

long_factor: Optional[float] = None

long_mscale: Optional[float] = None

low_freq_factor: Optional[float] = None

original_max_position_embeddings: Optional[int] = None

replace(**kwargs)

rope_type: str = 'default'

short_factor: Optional[float] = None

short_mscale: Optional[float] = None

to_dict() → Dict[str, Any]

Convert the config to a dictionary, excluding None values.

to_tuple()

values() → an object providing a view on D's values

class easydel.layers.rotary_embedding.RotaryEmbedding(*args: Any, **kwargs: Any)

Bases: Module

class easydel.layers.rotary_embedding.YaRNScalingRotaryEmbedding(*args: Any, **kwargs: Any)

Bases: RotaryEmbedding

RotaryEmbedding extended with YaRN method.

Credits to Peng et al. github.com/jquesnelle/yarn

easydel.layers.rotary_embedding.apply_basic_rope(query: ~jax.Array, key: ~jax.Array, positions: ~jax.Array, frequencies: ~jax.Array, rotary_dim: int, is_neox_style: bool, offsets: ~jax.Array = None, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>)

easydel.layers.rotary_embedding.apply_phi3_rope(query, key, positions, frequencies, offsets: ~jax.Array = None, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>)

easydel.layers.rotary_embedding.compute_basic_frequencies(base: int, rotary_dim: int, max_position_embeddings: int)

easydel.layers.rotary_embedding.compute_basic_inv_frequencies(base: int, rotary_dim: int)

easydel.layers.rotary_embedding.compute_deepseek_frequencies(base, rotary_dim, scaling_factor, extrapolation_factor, beta_fast, beta_slow, max_position_embeddings, mscale, mscale_all_dim, attn_factor) → Array

easydel.layers.rotary_embedding.compute_dynamic_frequencies(base: int, rotary_dim: int, max_position_embeddings: int, scaling_factor: float)

easydel.layers.rotary_embedding.compute_linear_frequencies(base: int, rotary_dim: int, max_position_embeddings: int, scaling_factors: List[float])

easydel.layers.rotary_embedding.compute_llama3_frequencies(base, rotary_dim, low_freq_factor, high_freq_factor, scaling_factor, max_position_embeddings: int)

easydel.layers.rotary_embedding.compute_llama3_inv_frequencies(base, rotary_dim, low_freq_factor, high_freq_factor, orig_max_position, scaling_factor)

easydel.layers.rotary_embedding.compute_phi3_frequencies(base, head_size, rotary_dim, max_position_embeddings, original_max_position_embeddings, short_factor, long_factor)

easydel.layers.rotary_embedding.compute_yarn_frequencies(base: float, rotary_dim: int, beta_fast: float, beta_slow: float, max_position_embeddings: int, scaling_factor: float, extrapolation_factor: float, attn_factor: float) → Array

easydel.layers.rotary_embedding.compute_yarn_inv_frequencies(base: float, rotary_dim: int, beta_fast: float, beta_slow: float, max_position_embeddings: int, scaling_factor: float, extrapolation_factor: float) → Array

easydel.layers.rotary_embedding.get_frequencies(head_size: int, rotary_dim: int, max_position: int, base: int, rope_scaling: Optional[Dict[str, Any]] = None, partial_rotary_factor: float = 1.0) → Array

easydel.layers.rotary_embedding.get_rope(head_size: int, rotary_dim: int, max_position: int, base: int, is_neox_style: bool = True, rope_scaling: Optional[Dict[str, Any]] = None, dtype: Optional[dtype] = None, partial_rotary_factor: float = 1.0) → RotaryEmbedding

easydel.layers.rotary_embedding.rope_wraper(type)

easydel.layers.rotary_embedding.yarn_get_mscale(scale: float = 1, mscale: float = 1) → float