def METHOD_NAME(vstate_rho, normalize):
    rho = vstate_rho.to_matrix(normalize=normalize)
    if normalize:
        np.testing.assert_allclose(jnp.trace(rho), 1.0)
    rho_norm = rho / jnp.trace(rho)
    assert rho.shape == (
        vstate_rho.hilbert.physical.n_states,
        vstate_rho.hilbert.physical.n_states,
    )
    x = vstate_rho.hilbert.all_states()
    rho_exact = jnp.exp(vstate_rho.log_value(x)).reshape(rho.shape)
    rho_exact = rho_exact / jnp.trace(rho_exact)
    np.testing.assert_allclose(rho_norm, rho_exact)