mlx.core.fast.rope

Contents

mlx.core.fast.rope#

rope(a: array, dims: int, *, traditional: bool, base: float | None, scale: float, offset: int, freqs: array | None = None, stream: None | Stream | Device = None) → array#

Apply rotary positional encoding to the input.

Parameters:

  • a (array) – Input array.

  • dims (int) – The feature dimensions to be rotated. If the input feature is larger than dims then the rest is left unchanged.

  • traditional (bool) – If set to True choose the traditional implementation which rotates consecutive dimensions.

  • base (float, optional) – The base used to compute angular frequency for each dimension in the positional encodings. Exactly one of base and freqs must be None.

  • scale (float) – The scale used to scale the positions.

  • offset (int) – The position offset to start at.

  • freqs (array, optional) – Optional frequencies to use with RoPE. If set, the base parameter must be None. Default: None.

Returns:

The output array.

Return type:

array

Contents