Searched refs:set_sharding (Results 1 – 7 of 7) sorted by relevance
315 inst->set_sharding(operand_sharding); in HandleBatchNormTraining()317 inst->set_sharding(default_sharding); in HandleBatchNormTraining()320 tuple->set_sharding(sharding); in HandleBatchNormTraining()419 inst->set_sharding(sharding); in HandleBatchNormInference()421 inst->set_sharding(default_sharding); in HandleBatchNormInference()424 shifted_normalized->set_sharding(sharding); in HandleBatchNormInference()597 inst->set_sharding(activation_sharding); in HandleBatchNormGrad()599 inst->set_sharding(default_sharding); in HandleBatchNormGrad()602 tuple->set_sharding(sharding); in HandleBatchNormGrad()
127 gte->set_sharding(sharding); in FixupPassThroughDomainLinks()240 instruction->set_sharding(domain_sharding); in ApplyShardingFromUsers()260 instruction->set_sharding(domain_sharding); in ApplyShardingFromUsers()302 instruction->set_sharding(HloSharding::Tuple(sharding_tree)); in ApplyShardingFromUsers()305 instruction->set_sharding(sharding_tree.leaf_begin()->second); in ApplyShardingFromUsers()
151 p1->set_sharding(HloSharding::AssignDevice(1)); in TEST()162 p2->set_sharding(sharding); in TEST()
631 instruction->set_sharding(sharding); in CreateFromProto()1130 broadcast->set_sharding(operand->sharding()); in CreateBroadcastSequence()1155 reshaped_operand->set_sharding(operand->sharding()); in CreateBroadcastSequence()1162 broadcast->set_sharding(operand->sharding()); in CreateBroadcastSequence()1217 set_sharding(HloSharding::Tuple(sharding.GetAsShapeTree(shape()))); in set_single_sharding()1219 set_sharding(sharding); in set_single_sharding()1230 derived_instruction->set_sharding(*sharding_); in SetupDerivedInstruction()
1158 void set_sharding(const HloSharding& sharding) { in set_sharding() function1161 void set_sharding(std::shared_ptr<const HloSharding> sharding) { in set_sharding() function
871 copy->set_sharding(sharding); in SetupCopiedInstruction()
1722 instruction->set_sharding( in ParseInstructionRhs()