Searched refs:replicated_dimensions (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 1681 std::vector<int64> replicated_dimensions( in HandleSort() local
1683 replicated_dimensions[sort_dim] = per_partition_size * partition_count; in HandleSort()
1685 {ShapeUtil::MakeShape(element_type, replicated_dimensions), in HandleSort()
1686 ShapeUtil::MakeShape(index_type, replicated_dimensions)}); in HandleSort()
1704 replicated_dimensions[sort_dim] = k.value() * partition_count; in HandleSort()
1708 slice_input, ShapeUtil::MakeShape(element_type, replicated_dimensions), in HandleSort()
1718 slice_index, ShapeUtil::MakeShape(index_type, replicated_dimensions), in HandleSort()
1727 {ShapeUtil::MakeShape(element_type, replicated_dimensions), in HandleSort()
1728 ShapeUtil::MakeShape(index_type, replicated_dimensions)}); in HandleSort()
|