Gaussian Mixture Model

[krstopro]

Added Gaussian Mixture Model implementation, similar to the one in scikit-learn Python module. Currently it supports only full covariances and uses k-means for the initialisation. Other covariance types (diagonal, tied, spherical) and initialisation methods (random, k-means++) can easily be added.

[msluszniak]

Thank you for your contribution! The PR LGTM :)

[polvalente]

Lots of really minor comments. Overall this is looking great!

[polvalente]

From reading without looking at the code beforehand:

Does this mean that we need to run the algorithm multiple times, or does the algorithm itself repeat this procedure?

[krstopro]

The algorithm itself repeats the procedure num_runs amount of times, each time with different initial parameters. Similarly as it was already done in k-means implementation.

[krstopro]

I am not sure if I should change anything here. Thoughts?

[msluszniak]

For me, the description is clear. @polvalente do you have any suggestions?

[krstopro]

@msluszniak @polvalente Thanks for the support and given feedback! I've addressed most of the suggestions provided; see the comments above (except for vectorization that I would leave for now). There are few more remarks that I would like to make; see some of the comments I've written.

[msluszniak]

Are there any blockers for this PR?

[krstopro]

@msluszniak Don't think so. I just need to address few suggestions that you made. Could you have a look at the two reviews I opened? The logsumexp one might be important.

[msluszniak]

@msluszniak Don't think so. I just need to address few suggestions that you made. Could you have a look at the two reviews I opened? The logsumexp one might be important.

Sure, but tbh I cannot see them 😅

[krstopro]

Sure, but tbh I cannot see them 😅

@msluszniak They are probably collapsed among many others :)
Tagged you in the comments.

[msluszniak]

Are you sure that your comments are not in pending mode or in the batch and not pushed or sth? Because I really cannot find anything

[krstopro]

Are you sure that your comments are not in pending mode or in the batch and not pushed or sth? Because I really cannot find anything

Sorry, they were in pending mode and I don't think I am able to submit them. Hence, I'll write them here.
Line 357-366 includes implementation to logsumexp which basically calculates the logarithm of the sum of the exponentials of the tensor in a numerically stable way. This is a well known trick that is implemented in numerical libraries, e.g. see SciPy
https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.logsumexp.html.
However, I am not aware that such a function exists in Nx. Is there one or perhaps I should raise an issue about implementing it?

Line 292: doing {num_gaussians, num_features, num_features} = Nx.shape(covariances) inside defn won't work, even though it works outside. I suppose this is intended?

[msluszniak]

I'm almost sure that there is no such function implemented in Nx so feel free to send a PR with it. I'm not sure if we eventually want to have this type of functions in Nx or in Scholar, but it will be helpful.

This line won't work in defn because you're doing pattern matching. If you want to assert that the shape is as follows you need to check for it separately.

[polvalente]

You can use case for checking the validity of the shape inside defn

[krstopro]

You can use case for checking the validity of the shape inside defn

I see, thanks. Here it is not an issue, just wanted to check that this is intended.
Anyway, if everything else is fine can you approve the PR?

[polvalente]

Suggested change

	covariance = Nx.dot(responsibilities[[.., k]] * Nx.transpose(diff), diff) / nk[k]
	covariance = Nx.dot(responsibilities[[.., k]] * diff, [-2], diff / nk[k], [-2])

[krstopro]

This won't work as Nx.transpose is only applied to diff, not the entire product. Adding another axis to responsibilities[[.., k]] solves the problem.
covariance = Nx.dot(diff * Nx.new_axis(responsibilities[[.., k]], 1), [0], diff / nk[k], [0])

[josevalim]

@krstopro please verify @polvalente suggestions above and we can ship it :)

[josevalim]

💚 💙 💜 💛 ❤️