//  [python]
def test_dseigh_poly(
    rng_seeds,  # noqa: F811
    torch_devices,  # noqa: F811
    atol_poly,
    dims_ranks_outer_inner_posteriori,
    steep_poly_decay,
    recons_frob_tol,
    svectors_frob_tol,
    svectors_check_n,
):
    """Test distributed SEIGH on (PSD and non-PSD) polynomial matrices.

    Sample matrices, and compute their ``torch.eigh`` as well as the sketched
    EIGH. Test that:
    * Recovered eigenvalues are correct
    * Recovered matrix is correct and symmetric
    * The first bulk eigenvectors are correct
    * A-posteriori error estimation is low
    """
    for seed in rng_seeds:
        for dtype, atol in atol_poly.items():
            for d, r, (outer, inner), post in dims_ranks_outer_inner_posteriori:
                for dec in steep_poly_decay:
                    for psd in (True, False):
                        # we create matrix and full eigh on cpu, it is faster
                        mat = SynthMat.poly_decay(
                            shape=(d, d),
                            rank=r,
                            decay=dec,
                            symmetric=True,
                            seed=seed,
                            dtype=dtype,
                            device="cpu",
                            psd=psd,
                        )
                        S, U = torch.linalg.eigh(mat)  # mat = U @ diag(S) @ U.T
                        # sort eigdec in descending eigval magnitude:
                        _, perm = abs(S).sort(descending=True)
                        S, U = S[perm], U[:, perm]
                        # make sure we actually have negative eigvals
                        if not psd:
                            assert (S < -abs(atol)).any(), (
                                "non-PSD should have " + "negative eigvals!"
                            )
                        # then we test decomposition on both devices
                        for device in torch_devices:
                            mat = mat.to(device)
                            S = S.to(device)
                            U = U.to(device)
                            #
                            tmpdir = tempfile.TemporaryDirectory()
                            (Q, core_U, core_S) = mock_dseigh(
                                tmpdir.name,
                                mat,
                                device,
                                dtype,
                                outer,
                                inner,
                                seed,
                            )
                            # check that singular values are correct
                            assert torch.allclose(
                                S[: 2 * r], core_S[: 2 * r], atol=atol
                            ), "Bad recovery of singular values!"

                            # check that recovered matrix is the same
                            core_U, core_S = truncate_core(
                                2 * r, core_U, core_S
                            )
                            mat_recons = (
                                Q @ core_U @ torch.diag(core_S) @ core_U.T @ Q.T
                            )
                            # check that reconstruction is symmetric
                            assert torch.allclose(
                                mat_recons, mat_recons.T, atol=atol
                            ), "Recovered sseigh is not symmetric!"
                            recons_err = (
                                torch.dist(mat, mat_recons) / mat.norm()
                            )
                            assert recons_err <= abs(
                                recons_frob_tol[dtype]
                            ), f"Bad SSVD reconstruction! {recons_err}"

                            # check a few singular vectors are correct: avoid
                            # the ones that have the same singular values, since
                            # they aren't unique. And compare outer prods, since
                            # multiplying both sides by -1 yields same result.
                            recons_evs = Q @ core_U[:, r:]
                            for idx in range(svectors_check_n):
                                recons_err = torch.dist(
                                    recons_evs[:, idx], U[:, r + idx]
                                ).item()
                                recons_err_flip = torch.dist(
                                    recons_evs[:, idx], -U[:, r + idx]
                                ).item()
                                recons_err = min(recons_err, recons_err_flip)
                                assert recons_err <= abs(
                                    svectors_frob_tol[dtype]
                                ), "Bad recovery of singular vectors!"
                            # check a-posteriori error. With 30 measurements, we
                            # know that the prob of the actual error being
                            # twice as big is less than 1%, and keeps shrinking
                            # exponentially: a_posteriori_error_bounds(30, 1)
                            # Therefore, a very small estimate almost certainly
                            # guarantees that the approximation is good, and we
                            # have enough rank in our truncated SVD.
                            for adjoint in [True, False]:
                                _, _, err_estimate = a_posteriori_error(
                                    mat,
                                    mat_recons,
                                    post,
                                    seed=seed,
                                    dtype=dtype,
                                    device=device,
                                    adjoint=adjoint,
                                )[0]
                                assert (
                                    err_estimate / mat.numel()
                                ) <= atol, "Bad a-posteriori estimate!"
                            # cleanup HDF5 tmpdir
                            tmpdir.cleanup()