def jacobi_svd(A): U, D, V = house_bidiag(A) m, n = D.shape iter, max_iter = 0, 100 frobenius_norm = np.linalg.norm(D) diag_norm = np.linalg.norm(np.diag(D)) off_diag_norm = np.sqrt( frobenius_norm - diag_norm) * np.sqrt(frobenius_norm + diag_norm) while off_diag_norm > 1e-6 * frobenius_norm and iter < max_iter: iter += 1 for p in range(m - 1): for q in range(p + 1, n): rot_l, rot_r = jacobi_rot(D
| 806 | // return U, np.diag(D), V |
| 807 | // |
| 808 | SVDResult SVD(XlaOp a, int64 max_iter, float epsilon, |
| 809 | PrecisionConfig::Precision precision) { |
| 810 | XlaBuilder* builder = a.builder(); |
| 811 | auto return_error = [&](const Status& status) { |
| 812 | SVDResult result; |
| 813 | result.u = builder->ReportError(status); |
| 814 | result.v = builder->ReportError(status); |
| 815 | result.d = builder->ReportError(status); |
| 816 | return result; |
| 817 | }; |
| 818 | auto shape_with_status = builder->GetShape(a); |
| 819 | if (!shape_with_status.status().ok()) { |
| 820 | return return_error(shape_with_status.status()); |
| 821 | } |
| 822 | Shape a_shape = shape_with_status.ValueOrDie(); |
| 823 | const int64 num_dims = a_shape.rank(); |
| 824 | const int64 num_batch_dims = num_dims - 2; |
| 825 | std::vector<int64> batch_dims(num_batch_dims); |
| 826 | for (int i = 0; i < num_batch_dims; ++i) { |
| 827 | batch_dims[i] = ShapeUtil::GetDimension(a_shape, i); |
| 828 | } |
| 829 | int64 m = ShapeUtil::GetDimension(a_shape, -2); |
| 830 | int64 n = ShapeUtil::GetDimension(a_shape, -1); |
| 831 | bool maybe_transpose = m < n; |
| 832 | |
| 833 | if (maybe_transpose) { |
| 834 | a = TransposeInMinorDims(a); |
| 835 | std::swap(m, n); |
| 836 | } |
| 837 | |
| 838 | auto eps = ScalarLike(a, epsilon); |
| 839 | |
| 840 | SVDResult svd_result = |
| 841 | HouseHolderBidiagonalization(a, eps, precision).ValueOrDie(); |
| 842 | |
| 843 | auto output_with_status = WhileLoopFn( |
| 844 | { |
| 845 | Zero(builder, S32), // k |
| 846 | svd_result.u, // u |
| 847 | svd_result.v, // v |
| 848 | svd_result.d, // d |
| 849 | eps, // epsilon |
| 850 | }, // |
| 851 | n, // |
| 852 | max_iter, // |
| 853 | "CyclicOneSidedJacobi", // |
| 854 | builder); |
| 855 | if (!output_with_status.status().ok()) { |
| 856 | return return_error(output_with_status.status()); |
| 857 | } |
| 858 | |
| 859 | auto output = output_with_status.ValueOrDie(); |
| 860 | |
| 861 | svd_result.u = output[1]; |
| 862 | svd_result.v = output[2]; |
| 863 | svd_result.d = output[3]; |
| 864 | svd_result = SortBySingularValuesAndPostProcessing(svd_result).ValueOrDie(); |
| 865 | if (maybe_transpose) { |