RFC: add `materialize` to materialize lazy arrays

[lucascolley]

Preface

I do not think that I am the best person to champion this effort, as I am far from the most informed person here on Lazy arrays. I'm probably missing important things, but I would like to start this discussion as I think that it is an important topic.

The problem

The problem of mixing computation requiring data-dependent properties with lazy execution is discussed in detail elsewhere:

• https://data-apis.org/array-api/draft/design_topics/lazy_eager.html
• https://data-apis.org/array-api/draft/design_topics/data_dependent_output_shapes.html#data-dependent-output-shapes
• Handling materialization of lazy arrays #748
• Calculate number of unique values in a lazy array #834

A possible solution

Add the function materialize(x: Array) to the top level of the API. Behaviour:

• for eagerly-executed arrays, this would be a no-op
• for lazy arrays, this would force computation such that the data is available in the returned array (which is of the same array type?)
• for "100% lazy" arrays (Handling materialization of lazy arrays #748 (comment)), this would raise an exception

Prior art

• Dask:
  • https://docs.dask.org/en/stable/generated/dask.array.Array.compute_chunk_sizes.html computes chunk sizes / shape, working in-place and leaving the array as a Dask array.
  • https://docs.dask.org/en/stable/generated/dask.array.Array.compute.html materialises the in-memory equivalent of the dask array, returning e.g. a NumPy array.
• JAX:
  • ?
• others?

Concerns

• I think the main concern is whether eager-only libraries will agree to adding a no-op into the API. There is precedent for that type of change (e.g. device kwargs in NumPy), but perhaps this is too obtrusive?
• As far as I can tell there isn't a standard way to do this across lazy libraries. Does JAX just do this automatically when it would be needed? Do other libraries have this capability?

Alternatives

1. Do nothing. The easy option, but it leaves us unable to support lazy arrays when data-dependent properties are used in computation (maybe that is okay?)
2. An alternative API. Maybe spelled like compute* or a method on the array object. Maybe with options for partial materialization (if that's a thing)?

---

cc @TomNicholas @hameerabbasi @rgommers

[asmeurer]

A question is whether it's appropriate for an array API consuming library to materialize a lazy graph "behind the user's back", as it were. Such an operation could be quite expensive in general, and a user might be surprised to find that a seemingly innocuous function from something like scipy is doing this.

On the other hand, if an algorithm fundamentally depends on array values in a loop, there's no way it can be implemented without something like this. So maybe the answer is we should just provide guidance that functions that use materialize should clearly document that they do, or maybe even disable it unless given some sort of explicit flag by the user.

[asmeurer]

If I understand https://data-apis.org/array-api/draft/design_topics/lazy_eager.html correctly, the primary APIs that require materialization are __bool__, __int__, __float__, etc. https://data-apis.org/array-api/draft/design_topics/lazy_eager.html

So another option here would be to add a compute flag to a potential item API #815.

[hameerabbasi]

A couple of thoughts:

1. Regarding partial materialization, I don't think a separate API is necessary for that. Given there are already methods to select elements one wants from an array (indexing and so on), one can call materialize on the result of those. Whether or not that backfills the data of the indexed array is entirely the library's choice.
2. Regarding the "do nothing" option, it's kind of tricky -- sometimes an explicit barrier to compute the data is needed to "help" optimization of the compute graph.
3. Adding compute to the item API is something I'm -0.7 on if it's the only API. Ideally, you would compute only the part of the array that's needed, and doing it element by element in item would be excruciatingly slow in many cases.

[TomNicholas]

Thank you for starting this discussion @lucascolley, and thanks for tagging me!

function materialize(x: Array)

Note that the signature here should probably be more like materialize(*arrs: Array, **kwargs), as both dask.compute and cubed.compute can compute multiple lazy arrays at once. You often want to do that in order to handle common intermediate arrays efficiently. The **kwargs is also important because different parallel execution frameworks will require specifying different configuration options at runtime.

Prior art

Adding some more here:

• Cubed:
  • Cubed (cc @tomwhite) deliberately copies the lazy API of dask.array. It too has a compute function for multiple arrays (as well as a .compute method on single arrays). An intermediate Plan object is created which is analogous to dask's task graph, but consists entirely of lazy arrays (as opposed to arbitrary functions that can go in a dask task graph). Normally computed results are either plotted or serialized to Zarr on-disk.
• Xarray's "Lazy Indexing Classes"
  • Internally Xarray has an implementation of lazy indexing, which it uses to make sure that only data that is necessary for the result is loaded from on-disk files. Note this is completely separate from (and predates) dask's integration with xarray. It's not a fully-fledged lazy array API implementation, but there have been discussions about how it could become one (Lazy indexing arrays as a stand-alone package pydata/xarray#5081).
• JAX:
  • I'm not super familiar with JAX but if I understand correctly JAX is effectively lazy but distributed over "local XLA devices", which act a little bit like dask chunks. jax.pmap and jax.vmap can be used to distributed computation over the devices, and there are methods which trigger array computation e.g. .block_until_ready().

If I understand https://data-apis.org/array-api/draft/design_topics/lazy_eager.html correctly, the primary APIs that require materialization are bool, int, float, etc.

I don't think this is the primary API at all, it's just an interesting special case where the return type is out of our hands. The primary API is as @lucascolley says, a materialize/compute function. For example in Xarray we basically compute if:

• the user asks to see numeric .values,
• the user tries to .plot their array,
• the user wants to save their array to disk (e.g. as netCDF/Zarr).

3. Adding compute to the item API is something I'm -0.7 on if it's the only API. Ideally, you would compute only the part of the array that's needed, and doing it element by element in item would be excruciatingly slow in many cases.

I agree - in Xarray we very rarely use .item, instead we are normally computing whole arrays (which have often been lazily indexed beforehand, so we are just computing a subarray of the original array).

I also agree with the other 2 points @hameerabbasi just made.

Xarray has a new abstraction over dask, cubed (and maybe soon JAX) called a "ChunkManager", (though I think we will rename it to a ComputeManager to better reflect its responsibilities). This can be understood as a way to create lazy arrays (.from_array), a way to distribute computation over them (.apply_gufunc), and a way to trigger computation of them (.compute).

[lucascolley]

Thanks all for the comments! Just tagging @jakevdp also who can maybe shed some light on JAX.

[jakevdp]

I don't think materialize would be useful for JAX. In that case, it would either be a no-op (during eager execution) or an unconditional error (during traced execution). When JAX is tracing a function, there is no concrete array to materialize!

[hameerabbasi]

Note that the signature here should probably be more like materialize(*arrs: Array, **kwargs), as both dask.compute and cubed.compute can compute multiple lazy arrays at once. You often want to do that in order to handle common intermediate arrays efficiently.

I'm +1 on an API that allows simultaneous materialisation of multiple arrays, although I'd spell it slightly differently.

1. To preserve type safety and avoid special cases that require wrapping/unwrapping of tuples, I'd make the type of the first argument Array | Iterable[Array].
2. The supplied kwargs are likely to be implementation-specific. Since the standard is a least common denominator, I feel this won't make it in. This, of course, doesn't prevent implementations from adding their own versions.

With this in mind, the signature I'd propose is materialize(x: Array | Iterable[Array]) -> Array | tuple[Array, ...]

[jakevdp]

One thing I'm unclear on: what is the difference between materialized and non-materialized arrays in terms of the array API? What array API operations can you do on one, but not on the other?

[hameerabbasi]

One thing I'm unclear on: what is the difference between materialized and non-materialized arrays in terms of the array API? What array API operations can you do on one, but not on the other?

I feel this is more driven by use-cases and performance. Some of these are outlined in #748 (comment) and #728.

[lucascolley]

What array API operations can you do on one, but not on the other?

One we bumped into in SciPy is xp.unique_values (#834)

[asmeurer]

One thing I'm unclear on: what is the difference between materialized and non-materialized arrays in terms of the array API? What array API operations can you do on one, but not on the other?

As I mentioned above, __bool__, __int__, __float__, __complex__, and __index__ will not work on unmaterialized arrays, unless the library is willing to automatically materialize them.

Additionally, the APIs that have data-dependent shapes are unique_all(), unique_counts(), unique_inverse(), and unique_values(), nonzero(), and repeat() when the repeats argument is an array, as well as boolean array indexing. But a lazy library would not necessarily disallow these. The shape of an array cannot be computed statically when one of these operations is used, but the array API allows None in shapes (for instance, Dask). The issue in #834 is not so much from using unique_values but from trying to call int on the resulting None dimension.