xla

write _shard_map; refactor flash attention to support 5d inputs.

#8730

Merged