Function OneSidedJacobiUpdate

tensorflow/compiler/xla/client/lib/svd.cc:451–563

Apply one-sided Jacobi on elements at indices pp, pq, qp, qq.

// Apply one-sided Jacobi on elements at indices pp, pq, qp, qq.
StatusOr<SVDResult> OneSidedJacobiUpdate(SVDResult svd_result, XlaOp p, XlaOp q,
                                         XlaOp eps) {
  XlaOp u = svd_result.u;
  XlaOp v = svd_result.v;
  XlaOp d = svd_result.d;
  XlaBuilder* builder = d.builder();
  TF_ASSIGN_OR_RETURN(Shape d_shape, builder->GetShape(d));
  const int64 num_dims = d_shape.rank();
  const int64 num_batch_dims = num_dims - 2;
  std::vector<int64> batch_dims(num_batch_dims);
  for (int i = 0; i < num_batch_dims; ++i) {
    batch_dims[i] = ShapeUtil::GetDimension(d_shape, i);
  }
  const int64 m = ShapeUtil::GetDimension(d_shape, -2);
  const int64 n = ShapeUtil::GetDimension(d_shape, -1);

  TF_ASSIGN_OR_RETURN(OneSidedJacobiRotation onesided_jacobi,
                      GetOneSidedJacobiRotation(d, p, q, eps));

  auto zero = ScalarLike(p, 0);

  // Zero out a_{pq} explicitly.
  std::vector<int64> pq_dims(batch_dims.begin(), batch_dims.end());
  pq_dims.push_back(1);
  pq_dims.push_back(1);
  auto pq_zero = ScalarLike(d, 0.0);
  auto pq_zeros = Broadcast(pq_zero, pq_dims);

  std::vector<int64> broadcast_dims(batch_dims.size());
  std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0);
  broadcast_dims.push_back(num_dims - 1);

  // Apply Jacobi Rotation on the left.
  auto slice_p = DynamicSliceInMinorDims(d, {p, zero}, {1, n});
  auto slice_q = DynamicSliceInMinorDims(d, {q, zero}, {1, n});
  auto slice_p_new =
      onesided_jacobi.rot_l.c * slice_p - onesided_jacobi.rot_l.s * slice_q;
  auto slice_q_new =
      onesided_jacobi.rot_l.s * slice_p + onesided_jacobi.rot_l.c * slice_q;
  d = DynamicUpdateSliceInMinorDims(d, slice_p_new, {p, zero});
  d = DynamicUpdateSliceInMinorDims(d, slice_q_new, {q, zero});

  // Apply Jacobi Rotation on the right.
  slice_p = DynamicSliceInMinorDims(d, {zero, p}, {m, 1});
  slice_q = DynamicSliceInMinorDims(d, {zero, q}, {m, 1});
  slice_p_new =
      onesided_jacobi.rot_r.c * slice_p - onesided_jacobi.rot_r.s * slice_q;
  slice_q_new =
      onesided_jacobi.rot_r.s * slice_p + onesided_jacobi.rot_r.c * slice_q;
  d = DynamicUpdateSliceInMinorDims(d, slice_p_new, {zero, p});
  d = DynamicUpdateSliceInMinorDims(d, slice_q_new, {zero, q});

  d = DynamicUpdateSliceInMinorDims(d, pq_zeros, {p, q});
  d = DynamicUpdateSliceInMinorDims(d, pq_zeros, {q, p});

  // Apply left Jacobi Rotation on U.
  slice_p = DynamicSliceInMinorDims(u, {zero, p}, {m, 1});
  slice_q = DynamicSliceInMinorDims(u, {zero, q}, {m, 1});