I think Inarritu having Stone look up, as if Keaton is flying, is his attempt to allow people to see a certain ambiguity-- the aforementioned magic realism-- that he undid by establishing firmly that every surreal thing he did/saw wasn't real. Which then leads me to wonder why the tone of the end of the film is so uplifting.