Searched defs:operand_sharding (Results 1 – 6 of 6) sorted by relevance
392 const HloSharding& operand_sharding = operand->sharding(); in InferDotShardingFromOperands() local460 const HloSharding& operand_sharding = operand->sharding(); in InferGatherParallelShardingFromOperands() local1037 const HloSharding& operand_sharding = annotate_op->operand(0)->sharding(); in InferUnspecifiedDimsFromOperand() local
274 HloSharding operand_sharding = in HandleBatchNormTraining() local
909 const Shape& operand_shape, const HloSharding& operand_sharding, in PassthroughOperandToGatherOutputOrScatterUpdate()
47 HloSharding operand_sharding; member
504 HloSharding operand_sharding = gather_sharding->operand_sharding; in PartitionGatherIndexParallelDimensions() local
213 if (auto operand_sharding = GetXlaShardingFromOperand(operand)) { in IdentifyXlaShardingForComputationInputs() local