Searched refs:temp_target_sharding (Results 1 – 1 of 1) sorted by relevance
998 const auto& temp_target_sharding = target_compatible_sharding.value(); in ReshardFromPartialReplicateWithDynamicSlice() local1000 if (temp_target_sharding.tile_assignment().dim(dim) > in ReshardFromPartialReplicateWithDynamicSlice()1005 temp_target_sharding.tile_assignment().dim(dim) / in ReshardFromPartialReplicateWithDynamicSlice()1017 hlo_, base_shape_, sharding(), temp_target_sharding, expand_tile_dims, in ReshardFromPartialReplicateWithDynamicSlice()1024 auto shard_shape = MakePartitionedShape(base_shape_, temp_target_sharding); in ReshardFromPartialReplicateWithDynamicSlice()1031 temp_target_sharding.tile_assignment().dim(i)); in ReshardFromPartialReplicateWithDynamicSlice()1033 auto offsets = MakePartitionOffsets(padded_base_shape, temp_target_sharding, in ReshardFromPartialReplicateWithDynamicSlice()1043 slice->set_sharding(temp_target_sharding); in ReshardFromPartialReplicateWithDynamicSlice()1047 if (CanReshardWithCollectivePermute(temp_target_sharding, target)) { in ReshardFromPartialReplicateWithDynamicSlice()