xla

c4f27718 - Improve some typing annotations (#8369)