def rope_p3

bendtorch@1.3.0.0 · main.bend:2288

import Base

2288# one pair (a = x[i1], b = x[i2]) rotated in place; fwd picks the direction: (a cos - b sin, b cos + a sin) forward,
2289# (a cos + b sin, b cos - a sin) backward, in PyTorch's order of operations for each
2290def rope_p3(fwd: Bool, +i1: U32, +i2: U32, +a: F32, r2: Array<F32> & F32, rc: Array<F32> & F32, rs: Array<F32> & F32) -> RS3:
2291  match fwd r2 rc rs:
2292    case True{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2293      RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (F32.neg(b) * z : F32) : F32)), i2, ((b * c : F32) + (a * z : F32) : F32)), cs, sn}
2294    case False{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2295      RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (b * z : F32) : F32)), i2, ((b * c : F32) + F32.neg((a * z : F32)) : F32)), cs, sn}