def METHOD_NAME(psi, ua, ub):
    """Assert UA, UB map the maximally entangled basis ``psi`` to the magic basis."""
    uaub = np.kron(ua, ub)
    for i, j in zip(range(4), [0, 1, 3, 2]):
        final_state = np.dot(uaub, psi[:, i])
        target_state = magic_basis[:, j]
        fidelity = np.abs(np.dot(np.conj(target_state), final_state))
        np.testing.assert_allclose(fidelity, 1)
        # np.testing.assert_allclose(final_state, target_state, atol=1e-12)