Delay allocating `embedding` array in `UMAP.fit_transform`

[jcrist]

Currently UMAP.fit_transform allocates an embedding array here and passes this to the c++ layer as an output parameter.

The embeddings are of size n_rows * n_components, which may take up a non-negligible amount of device memory. These embeddings are not filled in until the final stage in the algorithm, but take up device memory for an entire run.

A different way of handling this would be to have the C++ layer allocate the device array and return it. The python layer could then wrap this as UnownedMemory.

This wouldn't reduce peak device memory usage in UMAP (that currently happens in the embedding step, when we'd need the embedding array anyway), but it would reduce memory usage during e.g. the knn graph building step, which would let users use a lower nnd_n_clusters value for large datasets.

[divyegala]

I think we want the C++ layer to allocate the memory and then transfer the ownership to the Python layer. One way to do this would be to return a std::unique_ptr<rmm::device_buffer>> and use this API to take over ownership.

[viclafargue]

@divyegala @jinsolp Have you started working on this yet? Otherwise, I can help there.

[jinsolp]

Thanks @viclafargue , nothing from my side yet!

[csadorf]

@viclafargue I believe we had discussed that this might not be a priority for this release. Are there maybe some other UMAP issues that are higher priority?

[viclafargue]

Sure, I can start by working on the graph thresholding and outlier issues. The C++ side of this issue is almost complete though, so finishing it up should be fairly quick.