Tweak to work with new AA implementation