Searched refs:annotate_op (Results 1 – 1 of 1) sorted by relevance
1027 bool InferUnspecifiedDimsFromOperand(HloInstruction* annotate_op, in InferUnspecifiedDimsFromOperand() argument1032 CHECK(annotate_op->IsCustomCall("Sharding") || in InferUnspecifiedDimsFromOperand()1033 annotate_op->opcode() == HloOpcode::kCopy); in InferUnspecifiedDimsFromOperand()1034 if (!IsSpatiallyPartitioned(annotate_op->operand(0))) { in InferUnspecifiedDimsFromOperand()1037 const HloSharding& operand_sharding = annotate_op->operand(0)->sharding(); in InferUnspecifiedDimsFromOperand()1042 if (annotate_op->user_count() == 1) { in InferUnspecifiedDimsFromOperand()1043 HloInstruction* user = annotate_op->users()[0]; in InferUnspecifiedDimsFromOperand()1067 HloSharding sharding = annotate_op->sharding(); in InferUnspecifiedDimsFromOperand()1072 annotate_op->set_sharding(sharding); in InferUnspecifiedDimsFromOperand()1076 HloSharding auto_sharding = annotate_op->sharding(); in InferUnspecifiedDimsFromOperand()[all …]