| 111 | } |
| 112 | |
| 113 | StatusOr<JacobiUpdate> Update(JacobiUpdate jacobi_update, XlaOp p, XlaOp q, |
| 114 | XlaOp tol, int64 n) { |
| 115 | XlaBuilder* builder = jacobi_update.w.builder(); |
| 116 | TF_ASSIGN_OR_RETURN(JacobiRotation schur, SymmetricShurDecomposition2x2( |
| 117 | jacobi_update.w, p, q, tol)); |
| 118 | |
| 119 | TF_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(jacobi_update.w)); |
| 120 | const std::vector<int64> batch_dims(w_shape.dimensions().begin(), |
| 121 | w_shape.dimensions().end() - 2); |
| 122 | const int64 num_dims = w_shape.rank(); |
| 123 | |
| 124 | auto zero = ScalarLike(p, 0); |
| 125 | |
| 126 | XlaOp c = schur.c; |
| 127 | XlaOp s = schur.s; |
| 128 | |
| 129 | auto slice_p = DynamicSliceInMinorDims(jacobi_update.w, {p, zero}, {1, n}); |
| 130 | auto slice_q = DynamicSliceInMinorDims(jacobi_update.w, {q, zero}, {1, n}); |
| 131 | |
| 132 | auto slice_p_new = c * slice_p - s * slice_q; |
| 133 | auto slice_q_new = s * slice_p + c * slice_q; |
| 134 | |
| 135 | jacobi_update.w = |
| 136 | DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_p_new, {p, zero}); |
| 137 | jacobi_update.w = |
| 138 | DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_q_new, {q, zero}); |
| 139 | |
| 140 | slice_p = DynamicSliceInMinorDims(jacobi_update.w, {zero, p}, {n, 1}); |
| 141 | slice_q = DynamicSliceInMinorDims(jacobi_update.w, {zero, q}, {n, 1}); |
| 142 | |
| 143 | slice_p_new = c * slice_p - s * slice_q; |
| 144 | slice_q_new = s * slice_p + c * slice_q; |
| 145 | |
| 146 | jacobi_update.w = |
| 147 | DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_p_new, {zero, p}); |
| 148 | jacobi_update.w = |
| 149 | DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_q_new, {zero, q}); |
| 150 | |
| 151 | // Zero out a_{pq} explicitly. |
| 152 | std::vector<int64> pq_dims(batch_dims.begin(), batch_dims.end()); |
| 153 | pq_dims.push_back(1); |
| 154 | pq_dims.push_back(1); |
| 155 | auto pq_zero = ScalarLike(jacobi_update.w, 0.0); |
| 156 | auto pq_zeros = Broadcast(pq_zero, pq_dims); |
| 157 | jacobi_update.w = |
| 158 | DynamicUpdateSliceInMinorDims(jacobi_update.w, pq_zeros, {p, q}); |
| 159 | jacobi_update.w = |
| 160 | DynamicUpdateSliceInMinorDims(jacobi_update.w, pq_zeros, {q, p}); |
| 161 | |
| 162 | slice_p = DynamicSliceInMinorDims(jacobi_update.v, {zero, p}, {n, 1}); |
| 163 | slice_q = DynamicSliceInMinorDims(jacobi_update.v, {zero, q}, {n, 1}); |
| 164 | |
| 165 | std::vector<int64> broadcast_dims(batch_dims.size()); |
| 166 | std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); |
| 167 | broadcast_dims.push_back(num_dims - 1); |
| 168 | |
| 169 | // Renormalize the p-th and q-th columns. This step is redundant if high |
| 170 | // precision floats are used, like 64-bit float. But for 32-bit float, it |
no test coverage detected