Searched refs:ReplicateOnLastTileDim (Results 1 – 8 of 8) sorted by relevance
81 if (!may_combine_partial_sharding || !old.ReplicateOnLastTileDim() || in MergeSharding()82 !to_merge->ReplicateOnLastTileDim() || in MergeSharding()274 if (sharding.ReplicateOnLastTileDim() && in TransposeSharding()292 return sharding.ReplicateOnLastTileDim() in TransposeSharding()401 if (sharding.ReplicateOnLastTileDim()) { in ReshapeSharding()406 return sharding.ReplicateOnLastTileDim() in ReshapeSharding()427 return sharding.ReplicateOnLastTileDim() in ReverseSharding()519 if (index_sharding.ReplicateOnLastTileDim()) { in GatherOutputSharding()530 return index_sharding.ReplicateOnLastTileDim() in GatherOutputSharding()561 if (output_sharding.ReplicateOnLastTileDim()) { in GatherIndexSharding()[all …]
295 CHECK_EQ(shape.dimensions_size() + (ReplicateOnLastTileDim() ? 1 : 0), in TileLimitForDevice()562 result.set_replicate_on_last_tile_dim(ReplicateOnLastTileDim()); in ToProto()603 if (ReplicateOnLastTileDim()) { in NumTiles()615 CHECK(!ReplicateOnLastTileDim() || in NumTiles()
439 if (replicate_non_parallel_dims.ReplicateOnLastTileDim()) { in InferGatherParallelShardingFromOperands()445 return replicate_non_parallel_dims.ReplicateOnLastTileDim() in InferGatherParallelShardingFromOperands()761 if (op->sharding().ReplicateOnLastTileDim()) { in InferShardingFromOperands()768 op->sharding().ReplicateOnLastTileDim() in InferShardingFromOperands()1323 (user_sharding.ReplicateOnLastTileDim() ? 1 : 0)); in GetShardingFromUser()1336 return user_sharding.ReplicateOnLastTileDim() in GetShardingFromUser()
156 bool ReplicateOnLastTileDim() const { return replicate_on_last_tile_dim_; } in ReplicateOnLastTileDim() function
243 if (sharding.ReplicateOnLastTileDim()) { in MakeTiledPartitionOrdinals()305 if (!partial_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()310 (target_sharding.ReplicateOnLastTileDim() ? 1 : 0); in PartialReplicateReshardCompatibleSharding()349 if (target_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()364 if (target_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()381 target_sharding.ReplicateOnLastTileDim() in PartialReplicateReshardCompatibleSharding()403 return target_sharding.ReplicateOnLastTileDim() in PartialReplicateReshardCompatibleSharding()634 int64 rank = sharding.ReplicateOnLastTileDim() in UniqueTiledDim()1374 source.ReplicateOnLastTileDim() == target.ReplicateOnLastTileDim() && in CanReshardWithCollectivePermute()1414 if (sharding.ReplicateOnLastTileDim()) { in GroupShardingOnDims()[all …]
287 if (operand.sharding().ReplicateOnLastTileDim()) { in ParititonTrivialIndexedOperandDimension()358 if (indices_sharding.ReplicateOnLastTileDim()) { in PartitionIndexParallelDimensions()366 indices_sharding.ReplicateOnLastTileDim() in PartitionIndexParallelDimensions()580 if (indices.sharding().ReplicateOnLastTileDim()) { in HandleScatter()
418 if (!target.IsTileMaximal() && sharding().ReplicateOnLastTileDim()) { in ReshardNoCache()429 if (!sharding().IsTileMaximal() && target.ReplicateOnLastTileDim()) { in ReshardNoCache()455 if (target.ReplicateOnLastTileDim()) { in ReshardNoCache()912 if (!target.ReplicateOnLastTileDim()) { in ReshardToPartialReplicateWithAllGather()978 if (!sharding().ReplicateOnLastTileDim()) { in ReshardFromPartialReplicateWithDynamicSlice()1010 if (target.ReplicateOnLastTileDim()) { in ReshardFromPartialReplicateWithDynamicSlice()1139 auto temp_target = target.ReplicateOnLastTileDim() in ReshardWithAllToAll()1229 bool source_is_partial_replicate = sharding().ReplicateOnLastTileDim(); in ReshardPartialReplicateWithAllToAll()1233 if (!partial_replicate_sharding.ReplicateOnLastTileDim()) { in ReshardPartialReplicateWithAllToAll()1239 if (tile_sharding.ReplicateOnLastTileDim() || tile_sharding.IsTileMaximal()) { in ReshardPartialReplicateWithAllToAll()[all …]
486 if (lhs_sharding.ReplicateOnLastTileDim() || in PartitionBaseCase()487 rhs_sharding.ReplicateOnLastTileDim() || in PartitionBaseCase()488 output_sharding.ReplicateOnLastTileDim()) { in PartitionBaseCase()1722 output_sharding.ReplicateOnLastTileDim() in PartitionDotGroupOnBatch()1777 if (operand.sharding().ReplicateOnLastTileDim() && in PartitionDotGroupOnBatch()1903 matching_sharding.ReplicateOnLastTileDim() in GetNonContractingPartitionGroupedShardingForMatchedOperand()1930 if (other_sharding.ReplicateOnLastTileDim() && in GetNonContractingPartitionGroupedShardingForOtherOperand()2117 rhs_sharding = rhs_sharding.ReplicateOnLastTileDim() in PartitionDotGroupOnContracting()2126 lhs_sharding = lhs_sharding.ReplicateOnLastTileDim() in PartitionDotGroupOnContracting()2176 if (output_sharding.ReplicateOnLastTileDim() && in PartitionDotGroupOnContracting()[all …]