One sided Jacobi rotations. For a matrix, [a_pp, a_pq] [a_qp, a_qq] After applying Jacobi rotations on both sides, the matrix is diagonalized. [b_pp, 0] [0, b_qq] def jacobi_rot(a, p, q, eps): t = a[p, p] + a[q, q] d = a[q, p] - a[p, q] if np.abs(d) < eps: s = 0.0 c = 1.0 else: u = t / d tmp = np.sqrt(1.0 + u**2) s = -1.0 / tmp c = u / tmp rot = np.array([[c, s], [-s, c]]) m_tmp = rot.T @ a[[p,
| 411 | // return rot_l, rot_r |
| 412 | // |
| 413 | StatusOr<OneSidedJacobiRotation> GetOneSidedJacobiRotation(XlaOp a, XlaOp p, |
| 414 | XlaOp q, XlaOp eps) { |
| 415 | XlaOp a_pp = DynamicSliceInMinorDims(a, {p, p}, {1, 1}); |
| 416 | XlaOp a_pq = DynamicSliceInMinorDims(a, {p, q}, {1, 1}); |
| 417 | XlaOp a_qp = DynamicSliceInMinorDims(a, {q, p}, {1, 1}); |
| 418 | XlaOp a_qq = DynamicSliceInMinorDims(a, {q, q}, {1, 1}); |
| 419 | |
| 420 | XlaOp one = ScalarLike(a, 1.0); |
| 421 | |
| 422 | XlaOp t = a_pp + a_qq; |
| 423 | XlaOp d = a_qp - a_pq; |
| 424 | |
| 425 | XlaOp u = Div(t, d); |
| 426 | XlaOp tmp = Rsqrt(one + Square(u)); |
| 427 | |
| 428 | JacobiRotation rot; |
| 429 | |
| 430 | XlaOp zeros = ZerosLike(tmp); |
| 431 | XlaOp ones = zeros + one; |
| 432 | |
| 433 | rot.s = Select(Lt(Abs(d), eps), zeros, -tmp); |
| 434 | rot.c = Select(Lt(Abs(d), eps), ones, Mul(u, tmp)); |
| 435 | |
| 436 | XlaOp a_pp_new = rot.c * a_pp - rot.s * a_qp; |
| 437 | XlaOp a_pq_new = rot.c * a_pq - rot.s * a_qq; |
| 438 | XlaOp a_qq_new = rot.s * a_pq + rot.c * a_qq; |
| 439 | |
| 440 | OneSidedJacobiRotation rots; |
| 441 | TF_ASSIGN_OR_RETURN(rots.rot_r, |
| 442 | MakeJacobi(a_pp_new, a_qq_new, a_pq_new, eps)); |
| 443 | |
| 444 | rots.rot_l.c = rot.c * rots.rot_r.c - rot.s * rots.rot_r.s; |
| 445 | rots.rot_l.s = rot.s * rots.rot_r.c + rot.c * rots.rot_r.s; |
| 446 | |
| 447 | return rots; |
| 448 | } |
| 449 | |
| 450 | // Apply one-sided Jacobi on elements at indices pp, pq, qp, qq. |
| 451 | StatusOr<SVDResult> OneSidedJacobiUpdate(SVDResult svd_result, XlaOp p, XlaOp q, |
nothing calls this directly
no test coverage detected