Function MakeJacobi

tensorflow/compiler/xla/client/lib/svd.cc:359–384

MakeJacobi computes a rotation matrix G = [[c, s], [-s, c]], such that G_T * [[ps, pqs], [pqs, qs]] * G is diagonalized. def make_jacobi(ps, qs, pqs, eps): if np.abs(a_pq) > eps: tau = (a_qq - a_pp) / (2 * a_pq) if tau >= 0: t = 1.0 / (tau + np.sqrt(1 + tau ** 2)) else: t = -1.0 / (-tau + np.sqrt(1 + tau ** 2)) c = 1.0 / np.sqrt(1.0 + t ** 2) s = t * c else: c = 1.0 s = 0.0 return c, s

//     return c, s
//
StatusOr<JacobiRotation> MakeJacobi(XlaOp ps, XlaOp qs, XlaOp pqs, XlaOp eps) {
  auto zero = ScalarLike(ps, 0.0);
  auto one = ScalarLike(ps, 1.0);
  auto two = ScalarLike(ps, 2.0);

  auto tau = (qs - ps) / (pqs * two);
  auto t_pos = one / (tau + Sqrt(one + Square(tau)));
  auto t_neg = -one / (-tau + Sqrt(one + Square(tau)));
  auto t = Select(Ge(tau, zero), t_pos, t_neg);

  auto c_temp = Rsqrt(one + Square(t));
  auto s_temp = t * c_temp;

  auto c = Select(Ge(Abs(pqs), eps), c_temp, ZerosLike(c_temp) + one);
  auto s = Select(Ge(Abs(pqs), eps), s_temp, ZerosLike(s_temp));
  // Renormalize c and s to compensate for low precision arithmetic, this step
  // is redundant if high precision float is used, like float64.
  auto rnorm = Rsqrt(Square(c) + Square(s));

  JacobiRotation rot;

  rot.c = c * rnorm;
  rot.s = s * rnorm;

  return rot;
}

// One sided Jacobi rotations. For a matrix,
//  [a_pp, a_pq]