Searched refs:tile_dims (Results 1 – 2 of 2) sorted by relevance
449 std::vector<int64> tile_dims(sharding.tile_assignment().num_dimensions(), 1); in ReshapeToTileDimension() local457 tile_dims[i] = size; in ReshapeToTileDimension()481 tile_dims[dim] = devices.size() / ignore_size; in ReshapeToTileDimension()482 Array<int64> tile_assignment(tile_dims); in ReshapeToTileDimension()
1804 absl::InlinedVector<int64, 4> tile_dims(target_shape.dimensions_size(), 1); in CreateMatchingShardingOnDims() local1808 tile_dims[target_dims[i]] = in CreateMatchingShardingOnDims()1817 tile_dims.push_back(source_sharding.tile_assignment().num_elements() / in CreateMatchingShardingOnDims()1821 tgt_tile_assignment.Reshape(tile_dims); in CreateMatchingShardingOnDims()