Searched refs:new_tile_dims (Results 1 – 2 of 2) sorted by relevance
55 auto new_tile_dims = group_tile_assignment.dimensions(); in PartialTile() local56 new_tile_dims.push_back(replication_groups[0].size()); in PartialTile()57 auto new_tile_assignment = Array<int64>(new_tile_dims); in PartialTile()76 auto new_tile_dims = tile_assignment_last_dim_replicate.dimensions(); in PartialTile() local77 new_tile_dims.pop_back(); in PartialTile()79 fully_tiled.Reshape(new_tile_dims); in PartialTile()
89 std::vector<int64> new_tile_dims; in MergeSharding() local91 new_tile_dims.reserve(to_merge->tile_assignment().num_dimensions()); in MergeSharding()96 new_tile_dims.push_back(old_dim); in MergeSharding()98 new_tile_dims.push_back(new_dim); in MergeSharding()100 new_tile_dims.push_back(new_dim); in MergeSharding()106 int64 replication = num_devices / Product(new_tile_dims); in MergeSharding()107 if (!compatible || num_devices % Product(new_tile_dims) != 0 || in MergeSharding()111 new_tile_dims.push_back(replication); in MergeSharding()112 Array<int64> new_tile(new_tile_dims); in MergeSharding()186 new_tile_dims.pop_back(); in MergeSharding()[all …]