I'm planning on extracting the calls here to a compute function to realize lazy arrays.

Not sure if it will live in scikit-learn, though. I think I'd want it to live in array-api-compat, but I haven't discussed this with the folks there yet.

Based on usage here, the new compute function could have an option to e.g. compute shape only if that's what's needed, or compute the full array.

I think this kind of code needs to go somewhere else. Having the "lazyness" leak through is a bit silly. I'm not sure where it should be though.

I don't know how to resolve this issue. In the past I've suggested that accessing something like .shape should just trigger the computation for you or that everything needs to grow a .compute, even eager libs. Both would allow array consumers like scikit-learn to write code that does not care whether the implementation is lazy or not. However there doesn't seem to be much support for this within the Array API community. The alternative is having to place code that checks "is this lazy? if yes trigger compute" in consumers like scikit-learn which I think is not great.

So this is maybe not a task for you alone to solve @lithomas1 but we need to find some kind of solution

Thinking about this some more, would scikit-learn be happy with a shape helper (as opposed to my earlier suggestion of a compute helper) from array-api-compat?

This would handle the laziness of the various array libraries out there and call compute (or whatever the equivalent is) on the array if the shape needs to be materialized, before returning the shape of the array.

(In the event that array-api-compat is not installed, we can just define this to return, e.g. x.shape like scikit-learn already does)

The advantage of this approach would be that scikit-learn doesn't have to think about laziness anymore - that work would be outsourced to array-api-compat.

The only downsides of this approach of this approach are that

• We indiscriminately materialize (for lazy arrays) even if it's strictly not necessary (e.g. we just access the shape to pass it to an array constructor like np.zeros). I don't think we'll lose too much (if any) performance here, though.
• scikit-learn devs need to remember that accessing .shape is banned, and existing usages have to be migrated.
  • This can be mitigated with a pre-commit hook, to automatically detect .shape accesses.
    I think with something like ruff, one can write a custom rule to automatically rewrite .shape accesses to shape(x)
  • For existing usages, this is something that can be done as part of the Array API migration process, so it shouldn't cause too much churn on its own

How does this sound to you?

@betatim any thoughts on the above?

Replying here and linking @ogrisel's comment #28588 (comment)

I think having a helper like shape() is the only(?) way forward for now. I'd not add it to array-api-compat but instead add it to scikit-learn in utils/_array_api.py - we already have a few helpers there for things that feel "scikit-learn specific".

More adventurous: I wonder if we can even wrap the dask namespace (and via that its arrays) to make it so what .shape access triggers the computation. That way people who edit scikit-learn's code base don't need to know anything about this issue.

Note that a compute helper could also help deal with the boolean-masked assignment problem in r2_score described in more details in this comment: #28588 (comment)

This also a lazy evaluation problem but not related to shape values.

I think having a helper like shape() is the only(?) way forward for now. I'd not add it to array-api-compat but instead add it to scikit-learn in utils/_array_api.py - we already have a few helpers there for things that feel "scikit-learn specific".

We could also do both: have a public helper in array-api-compat and a private scikit-learn specific helper in scikit-learn, that does nothing for libraries that are not accessed via array-api-compat (e.g. array-api-strict) as long the spec does not provide a standard way to deal with this.

More adventurous: I wonder if we can even wrap the dask namespace (and via that its arrays) to make it so what .shape access triggers the computation. That way people who edit scikit-learn's code base don't need to know anything about this issue.

Not sure how feasible this is and whether its desirable or not to trigger computation implicitly when using lazy libraries.

To shed more light on what this is, I think this bug happens when scikit-learn calls unique_inverse.

I think something is going wrong in dask somewhere where the results of intermediate operations are getting corrupted.

When the error occurs in computation is somewhat flaky, but it happens more often than not without the compute here.

Did you investigate the cause of the corruption? It might be worth reporting a minimal reproducer upstream.

> does not have precedence of over [], right? Trying this pattern locally yields:

TypeError: 'float' object is not subscriptable

I think you meant the following instead:

But then the following s = s[:, None] seems redundant and similarly for idx[:, None] in the call to where.

Also we should probably rename idx to something more correct and explicit such as strictly_positive_mask.

Thanks! I think I had this working locally and passing tests, but I must have messed something up on the merge.

I'm planning on circling back to this towards the end of summer, so feel free to take over this if you're interested in the meantime.

Not holding my breath, but also hoping that dask.array support improves in the meantime as well.

(The recent Array API updates suggest that sort is the most pressing thing that's missing in dask. Linalg wise, the eig family of methods is probably the next biggest missing feature in dask, we just haven't seen it come up yet since not a lot of estimators have been ported yet).