xla

cee0820e - write _shard_map; refactor flash attention to support 5d inputs. (#8730)