315990901cdcdca241013ee5e6c581391bae2523,geomstats/geometry/spd_matrices.py,SPDMetricBuresWasserstein,exp,#SPDMetricBuresWasserstein#Any#Any#,734

Before Change

        eigvals, eigvecs = gs.linalg.eigh(base_point)
        transp_eigvecs = gs.einsum("...ij->...ji", eigvecs)
        rotated_tangent_vec = gs.matmul(transp_eigvecs, tangent_vec)
        rotated_tangent_vec = gs.matmul(rotated_tangent_vec, eigvecs)

        ones = gs.ones(eigvals.shape)
        vertical_index = gs.einsum("...i,...j->...ij", eigvals, ones)
        horizontal_index = gs.einsum("...j,...i->...ij", eigvals, ones)
        coefficients = 1 / (vertical_index + horizontal_index)

        rotated_lyapnunov = gs.einsum("...ij,...ij->...ij",
                                      rotated_tangent_vec, coefficients)
        rotated_hessian = gs.einsum("...ij,...j,...jk->...ik",
                                    rotated_lyapnunov, eigvals,
                                    rotated_lyapnunov)
        hessian = gs.matmul(eigvecs, rotated_hessian)

After Change

            Riemannian exponential.
        
        eigvals, eigvecs = gs.linalg.eigh(base_point)
        transp_eigvecs = Matrices.transpose(eigvecs)
        rotated_tangent_vec = Matrices.mul(transp_eigvecs, tangent_vec,
                                           eigvecs)
        coefficients = 1 / (eigvals[..., :, None] + eigvals[..., None, :])
        rotated_sylvester = rotated_tangent_vec * coefficients
        rotated_hessian = gs.einsum("...ij,...j,...jk->...ik",
                                    rotated_sylvester, eigvals,
                                    rotated_sylvester)
        hessian = Matrices.mul(eigvecs, rotated_hessian, transp_eigvecs)

Instances