Searched refs:resharded (Results 1 – 2 of 2) sorted by relevance
359 auto resharded = ReshardNoCache(target); in Reshard() local360 state_.reshard_cache->per_hlo_cache[resharded.hlo()] in Reshard()363 cache.emplace_back(target, std::move(resharded)); in Reshard()366 return resharded; in Reshard()804 auto resharded = ExchangeHaloAndGetValidData( in ReshardAsWindowedInput() local811 if (!resharded) { in ReshardAsWindowedInput()816 visiting_hlo = *resharded; in ReshardAsWindowedInput()844 auto update_cache = [&](PartitionedHlo resharded) { in Replicate() argument845 state_.reshard_cache->per_hlo_cache[resharded.hlo()] in Replicate()848 cache.emplace_back(HloSharding::Replicate(), std::move(resharded)); in Replicate()[all …]
1817 auto resharded = operand.Reshard(UngroupSharding(grouped)); in PartitionDotGroupOnBatch() local1818 top_level_sharding_to_reset.emplace_back(resharded.hlo(), in PartitionDotGroupOnBatch()1819 resharded.sharding()); in PartitionDotGroupOnBatch()1820 resharded.hlo()->set_sharding(grouped.sharding); in PartitionDotGroupOnBatch()1821 return PartitionedHlo(resharded.hlo(), in PartitionDotGroupOnBatch()1845 if (auto resharded = reshard_to_output_batch( in PartitionDotGroupOnBatch() local1850 per_group_lhs = *resharded; in PartitionDotGroupOnBatch()1854 if (auto resharded = reshard_to_output_batch( in PartitionDotGroupOnBatch() local1859 per_group_rhs = *resharded; in PartitionDotGroupOnBatch()