ENH: Store fitted attributes as dpt and return dpt tensors

[fcharras]

The compatibility layer for having the sklearn test keep working is still TODO. Let alone that, I think this would be the final state of KMeansEngine.

[fcharras]

I wanted to duplicate the KMeansEngine class with KMeansEngineDebug class that would be loaded when engine sklearn_numba_dpex.debug is used rather than bare sklearn_numba_dpex. The difference between debug and vanilla would be that debug engine converts all fitted attributes and all outputs to numpy arrays. The usecase is using it in sklearn unit tests.

But with current plugin design in scikit-learn/scikit-learn#24497 I don't see a way around. One module can only implement one engine.

Maybe there's a workaround consisting in exposing several top level namespaces (e.g sklearn_numba_dpex and sklearn_numba_dpex_debug) in the same package using the packages list in setup.py ?

Or WDYT about changing the plugin spec to enable a different syntax that allow passing dotted string to the config_context ?

E.g rewriting https://github.com/scikit-learn/scikit-learn/pull/24497/files#diff-1d31de81e903bd6529fbe68f8009b7113e3b7de4f1465572ef88af4d03a7dc5bR35 such that entry points whose value prefixes match the user inputed string are selected. (i.e. using prefix selection rather than str.split(".", 1) to detect the name of the engine)

Edit: somewhat related to #21

[ogrisel]

I think we should allow to programmatically register additional provider names (without in addition the importlib.metadata registry) for testing / debug / benchmarking purposes.

from sklearn._engine import register_engine
register_engine(
    engine_name="kmeans",
    engine_class="sklearn_numba_dpex.kmeans.engine:KMeansEngineDebug",
    provider_name="sklearn_numba_dpex_debug",
)

In this case we break the fact that the provider name is usually the same as the toplevel import package of the engine class, so this might impose some refactoring of the engine lookup / import logic to avoid relying on this assumption.

[fcharras]

I agree although since it requires work on the sklearn side, I think we can still push this PR on with an emvironment variable based switch in the meantime, I'll add a commit going this direction

[fcharras]

Out of WIP

[jjerphan]

Thanks @fcharras. Here are a few comments.

[jjerphan]

Suggested change

	if sample_idx >= n_samples:
	# Early return if result is already False.
	if sample_idx >= n_samples or result[zero_idx] == zero_idx:

[fcharras]

Note that this add a read operation and doesn't stop future tasks, it just ensures that future tasks will exit early after the read. So on average it's better when clusterings are different but not when clusterings are the same. Still better since different clusterings should be the most common outcome.

[fcharras]

Comments should be addressed in last commit and suggestions applied.

[jjerphan]

LTGM. Thank you, @fcharras!

[fcharras]

So in fact make_get_nb_distinct_clusters_kernel does suffer from a race condition, the unittest sometimes output 6, sometimes 7, but it more seems like a bug in numba_dpex atomics rather than our code. Or maybe I'm not sure what "loaded atomically" should actually mean.

[fcharras]

The test was the issue, we have n_clusters = max(labels) + 1 but used n_clusters = max(labels).