xla

6e3992ac - allow disabling layout pinning in optimizer_step and reduce_gradients (#3556)