Searched refs:GetPartitionedHlo (Results 1 – 5 of 5) sorted by relevance
1395 new_operands.push_back(GetPartitionedHlo(operand).Reshard(sharding).hlo()); in DefaultAction()1446 GetPartitionedHlo(operand).hlo()->set_sharding(operand->sharding()); in Preprocess()1453 logger_->RegisterLogEntry(GetPartitionedHlo(hlo).hlo(), in Postprocess()1460 GetPartitionedHlo(hlo).hlo()->set_sharding(*visiting_hlo_sharding_); in Postprocess()1464 GetPartitionedHlo(operand).hlo()->set_sharding(operand->sharding()); in Postprocess()1476 GetPartitionedHlo(operand).Reshard(hlo->sharding()).hlo()); in HandleElementwise()1497 GetPartitionedHlo(operand).Reshard(sharding).hlo()); in HandleConcatenate()1537 auto spmd_operand = GetPartitionedHlo(operand).Reshard(sharding).hlo(); in HandleConcatenate()1585 auto operand = GetPartitionedHlo(hlo->operand(0)).Reshard(sharding); in HandleSlice()1671 auto partitioned_input = GetPartitionedHlo(input).PadWithValue( in HandleSort()[all …]
510 auto operand = GetPartitionedHlo(scatter->operand(0)); in HandleScatter()511 auto indices = GetPartitionedHlo(scatter->operand(1)); in HandleScatter()512 auto updates = GetPartitionedHlo(scatter->operand(2)); in HandleScatter()674 auto operand = GetPartitionedHlo(gather->operand(0)); in HandleGather()675 auto indices = GetPartitionedHlo(gather->operand(1)); in HandleGather()
371 GetPartitionedHlo(hlo->operand(0)) in HandleFft()
455 PartitionedHlo& GetPartitionedHlo(const HloInstruction* hlo) { in GetPartitionedHlo() function
2840 auto& lhs = GetPartitionedHlo(hlo->operand(0)); in HandleDotHelper()2841 auto& rhs = GetPartitionedHlo(hlo->operand(1)); in HandleDotHelper()