Searched refs:reshape_dimensions (Results 1 – 2 of 2) sorted by relevance
353 auto reshape_dimensions = partial_sharding.tile_assignment().dimensions(); in PartialReplicateReshardCompatibleSharding() local354 int64 num_replication = reshape_dimensions.back(); in PartialReplicateReshardCompatibleSharding()360 reshape_dimensions.pop_back(); in PartialReplicateReshardCompatibleSharding()361 reshape_dimensions.insert(reshape_dimensions.end(), expand_tile_sizes.begin(), in PartialReplicateReshardCompatibleSharding()365 reshape_dimensions.push_back(num_target_replication); in PartialReplicateReshardCompatibleSharding()369 reshape_tile_assignment.Reshape(reshape_dimensions); in PartialReplicateReshardCompatibleSharding()
2077 std::vector<int64> reshape_dimensions( in SplitSpaceHelper() local2081 reshape_dimensions[spatial_dimension_to_split] = spatial_split_size; in SplitSpaceHelper()2082 reshape_dimensions[activations_batch_dim] = num_splits * old_batch_size; in SplitSpaceHelper()2085 MakeReshapeHlo(reshape_dimensions, activations)); in SplitSpaceHelper()