✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 8fd0b71. Link to the linter CI: here}

@ogrisel I added a small adjustment and the float32 tests seem to pass now. Do we have any way to invoke them in the CI as well?
Edit: I also updated two tests and adjusted the tolerance which checks for normalized weights because in the test test_gaussian_mixture_fit_predict it was failing due to the weights summing to around 1.0000001

Do we have any way to invoke them in the CI as well?

Yes, by pushing a commit with the [float32] marker in the commit message:

https://scikit-learn.org/stable/developers/contributing.html#commit-message-markers

Yes, by pushing a commit with the [float32] marker in the commit message:

https://scikit-learn.org/stable/developers/contributing.html#commit-message-markers

Okay nice let me see if the CI passes first before trying out float32 because for the initial pipelines the np.log did not accept a dtype argument.

I confirm that the tests now pass locally. Let me add a changelog entry.

Let me do a quick benchmark to see the impact of this change.

On my Apple M1 with accelerate (mamba install "libblas=*=*accelerate") I get the following numbers:

• on this branch: 7.8 s with a 1 GB memory usage peak
• on main: 23.1 s with a 2 GB memory usage peak

So a nice ~3X speed up and a 2X memory efficiency improvement.

Benchmark script:

from sklearn.mixture import GaussianMixture
from sklearn.datasets import make_blobs
import numpy as np
from time import perf_counter

X, y = make_blobs(n_samples=int(1e6), n_features=50, centers=10, random_state=0)
X = X.astype(np.float32)

tic = perf_counter()
GaussianMixture(n_components=10, covariance_type="full", random_state=0).fit(X)
toc = perf_counter()
print(f"Elapsed time: {toc - tic:.1f} s")

I used the scalene profiler to report memory usage (scalene bench_gmm.py with the script above to reproduce).

I also tried with OpenBLAS (mamba install "libblas=*=*openblas") and got a 1.9x speed up instead of nearly 3x. I suspect that Accelerate can use the Apple M1 GPU for some operations when working with float32 data while OpenBLAS is limited to the CPU vector instruction set.

Memory usage is the same for OpenBLAS and Accelerate.

I also tried with OpenBLAS (mamba install "libblas=*=*openblas") and got a 1.9x speed up instead of nearly 3x. I suspect that Accelerate can use the Apple M1 GPU for some operations when working with float32 data while OpenBLAS is limited to the CPU vector instruction set.

Even then there is a nice speedup.

Once this is merged, I think this estimator would be a nice candidate for array API support.

On my Apple M1 with accelerate (mamba install "libblas=*=*accelerate") I get the following numbers:

• on this branch: 7.8 s with a 1 GB memory usage peak
• on main: 23.1 s with a 2 GB memory usage peak

So a nice ~3X speed up and a 2X memory efficiency improvement.

Benchmark script:

from sklearn.mixture import GaussianMixture
from sklearn.datasets import make_blobs
import numpy as np
from time import perf_counter

X, y = make_blobs(n_samples=int(1e6), n_features=50, centers=10, random_state=0)
X = X.astype(np.float32)

tic = perf_counter()
GaussianMixture(n_components=10, covariance_type="full", random_state=0).fit(X)
toc = perf_counter()
print(f"Elapsed time: {toc - tic:.1f} s")

I used the scalene profiler to report memory usage (scalene bench_gmm.py with the script above to reproduce).

Great work! Beyond performance enhancement, in your tests, do you obtain identical results after downcasting X to np.float32?

Great work! Beyond performance enhancement, in your tests, do you obtain identical results after downcasting X to np.float32?

The existing tests still pass when we run them with float32 (up to some tolerance level that can depend on the choice of the dtype).

One question: is there a unit test that checks that the dtype of gmm.predict(X) matches the dtype of X or should it match the dtype of X that was used for fitting?

gmm.predict(X) returns the integer index of the component with the highest likelihood. So it is unrelated to the dtype of X.

Thanks, @lesteve and @OmarManzoor for the latest round of fixes / testing.