xla

cd73a62d - change all size node return value to be compatible with the underlying xla device. (#4410)