Searched refs:partition_id_ (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.h | 487 state.partition_id = partition_id_; in MakePartitioningState()
533 HloInstruction* partition_id_; variable
|
| D | spmd_partitioner.cc | 1363 partition_id_(collective_ops_creator_.create_partition_id(&b_)), in SpmdPartitioningVisitor()
1544 .Calculate(MakeTiledPartitionOrdinals(sharding, partition_id_, in HandleConcatenate()
1873 ShapeUtil::MakeShape(S32, partition_id_->shape().dimensions()), in HandleCustomCall()
1874 partition_id_)); in HandleCustomCall()
2156 MakeTiledPartitionOrdinals(sharding, partition_id_, &b_); in HandleIota()
2195 ShapeUtil::MakeShape(PRED, {}), partition_id_, on_device, in HandleSingleDevice()
2377 MakeTiledPartitionOrdinals(hlo->sharding(), partition_id_, &b_); in HandleDynamicUpdateSlice()
2396 ShapeUtil::ChangeElementType(partition_id_->shape(), PRED); in HandleDynamicUpdateSlice()
2547 branch_index = partition_id_; in HandleInfeed()
2559 ShapeUtil::MakeShape(S32, {1}), branch_index_table, {partition_id_}, in HandleInfeed()
[all …]
|
| D | gather_scatter_handler.cc | 579 auto partition_id = partition_id_; in HandleScatter()
650 operand, indices, partition_id_, scatter_dims_to_operand_dims, in HandleScatter()
|