Searched refs:halo_shape (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 833 auto halo_shape = hlo->shape(); in ExchangeHalo() local
836 halo_shape.set_dimensions(dim, halo_size); in ExchangeHalo()
837 std::vector<int64> halo_start_indices(halo_shape.rank(), 0); in ExchangeHalo()
839 std::vector<int64> halo_slice_strides(halo_shape.rank(), 1); in ExchangeHalo()
841 halo_shape, hlo, halo_start_indices, hlo->shape().dimensions(), in ExchangeHalo()
867 auto halo_shape = hlo->shape(); in ExchangeHalo() local
869 if (halo_size != halo_shape.dimensions(dim)) { in ExchangeHalo()
870 halo_shape.set_dimensions(dim, halo_size); in ExchangeHalo()
871 std::vector<int64> halo_start_indices(halo_shape.rank(), 0); in ExchangeHalo()
872 std::vector<int64> halo_slice_strides(halo_shape.rank(), 1); in ExchangeHalo()
[all …]
|