mirror of https://github.com/ollama/ollama
21 lines
660 B
Go
21 lines
660 B
Go
package nn
|
|
|
|
import (
|
|
"github.com/ollama/ollama/ml"
|
|
"github.com/ollama/ollama/ml/nn/rope"
|
|
)
|
|
|
|
// fastRoPE is an interface for tensors that support fast rotary positional embedding.
|
|
type fastRoPE interface {
|
|
RoPE(ctx ml.Context, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor
|
|
}
|
|
|
|
// RoPE applies rotary positional embedding to tensor `t`.
|
|
func RoPE(ctx ml.Context, t, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor {
|
|
if t, ok := t.(fastRoPE); ok {
|
|
return t.RoPE(ctx, positions, dim, base, scale, options...)
|
|
}
|
|
|
|
panic("RoPE not implemented for this tensor type")
|
|
}
|