28 lines
905 B
Swift
28 lines
905 B
Swift
// Copyright © 2024 Apple Inc.
|
|
|
|
import Foundation
|
|
import MLX
|
|
import MLXNN
|
|
|
|
// MARK: - applyRotaryPosition Helper
|
|
|
|
/// Apply rotary position embeddings, using the cache offset when available.
|
|
///
|
|
/// This function enables models to use a single call site instead of
|
|
/// repeating conditional offset handling:
|
|
/// ```swift
|
|
/// queries = applyRotaryPosition(rope, to: queries, cache: cache)
|
|
/// keys = applyRotaryPosition(rope, to: keys, cache: cache)
|
|
/// ```
|
|
///
|
|
/// - Parameters:
|
|
/// - rope: A RoPE layer conforming to both `OffsetLayer` and `ArrayOffsetLayer`.
|
|
/// - x: The input tensor to apply RoPE to.
|
|
/// - cache: The KV cache (determines offset), or `nil` for offset 0.
|
|
/// - Returns: The input with rotary positional encoding applied.
|
|
public func applyRotaryPosition<R: RoPELayer>(_ rope: R, to x: MLXArray, cache: KVCache?)
|
|
-> MLXArray
|
|
{
|
|
return rope(x, offset: cache?.offset ?? 0)
|
|
}
|