Searched refs:replicated_shape (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 1684 const Shape replicated_shape = ShapeUtil::MakeTupleShape( in HandleSort() local
1690 input_sharding.GetTupleSharding(replicated_shape).ValueOrDie(); in HandleSort()
1691 auto shard_shape = MakePartitionedShape(replicated_shape, topk_sharding); in HandleSort()
1839 const Shape replicated_shape = ShapeUtil::MakeTupleShape( in HandleCustomCall() local
1844 sharding.GetTupleSharding(replicated_shape).ValueOrDie(); in HandleCustomCall()
1846 MakePartitionedShape(replicated_shape, custom_call_sharding); in HandleCustomCall()
1851 PartitionedHlo partitioned_topk(topk, replicated_shape, in HandleCustomCall()
1913 replicated_shape, sort_dim, {replicated_value_gte, replicated_index_gte}, in HandleCustomCall()
1917 PartitionedHlo replicated_sort(sort, replicated_shape, in HandleCustomCall()
|