位置编码、RoPE 与 Transformer 模块

import torch
def half_layout_rope(x, position, base=10000.):
    # 教学片段前提：最后一维为正偶数，频率按每个二维旋转对生成
    half = x.shape[-1] // 2
    freq = base ** (-torch.arange(half, dtype=x.dtype) / half)
    angle = position[..., None] * freq  # 位置轴与频率轴形成角度表
    first, second = x[..., :half], x[..., half:]
    left = first * angle.cos() - second * angle.sin()
    right = second * angle.cos() + first * angle.sin()
    return torch.cat((left, right), -1)  # 前半与后半对应维度成对，而非相邻成对

public class Main {
  static double[] rotate(double[] x,double position,double base){
    if(x.length==0||x.length%2!=0||!Double.isFinite(position)||!Double.isFinite(base)||base<=1)
      throw new IllegalArgumentException();
    double[] out=new double[x.length];
    for(int i=0;i<x.length;i+=2){
      if(!Double.isFinite(x[i])||!Double.isFinite(x[i+1]))throw new IllegalArgumentException();
      double angle=position*Math.pow(base,-(double)i/x.length); // 第i/2对使用独立频率
      double c=Math.cos(angle),s=Math.sin(angle);
      out[i]=x[i]*c-x[i+1]*s;out[i+1]=x[i]*s+x[i+1]*c; // 相邻二维正交旋转
    }
    return out; // 保留输入数组，便于比较旋转前后性质
  }
  static double dot(double[] a,double[] b){double sum=0;for(int i=0;i<a.length;i++)sum+=a[i]*b[i];return sum;}
  static void close(double a,double b){if(Math.abs(a-b)>1e-9)throw new AssertionError(a+" != "+b);}
  public static void main(String[] args){
    double[] q={1,2,3,4},k={-2,1,0,5};
    close(rotate(new double[]{1,0},Math.PI/2,10000)[1],1);
    close(dot(rotate(q,3,10000),rotate(q,3,10000)),dot(q,q));
    close(dot(rotate(q,2,10000),rotate(k,5,10000)),
          dot(rotate(q,12,10000),rotate(k,15,10000)));
    double[] inverse=rotate(rotate(q,7,10000),-7,10000);
    for(int i=0;i<q.length;i++)close(inverse[i],q[i]);
    close(rotate(q,0,10000)[2],q[2]);
    boolean bad=false;try{rotate(new double[]{1,2,3},1,10000);}catch(IllegalArgumentException e){bad=true;}
    if(!bad)throw new AssertionError();
    System.out.println("6 rotary invariants passed");
  }
}