Stream: helpdesk (published)

Topic: Lazy array operations

view this post on Zulip Jesper Stemann Andersen (Dec 15 2024 at 10:58):

I'm working on MLX.jl to provide access to the MLX array framework (in particular for machine learning on Apple platforms).

In MLX almost all (array) operations are lazy (to make it possible to optimize a computational graph), i.e. operations are not executed until eval is called, e.g.:

>> import mlx.core as mx
>> a = mx.array([1, 2, 3, 4])
>> b = mx.array([1.0, 2.0, 3.0, 4.0])
>> c = a + b    # c not yet evaluated
>> mx.eval(c)  # evaluates c
>> c = a + b
>> print(c)     # Also evaluates c

(sorry for the "P-language" example :smile:)

How to best represent lazy arrays and operations on lazy arrays in Julia?
I have so far found Lazy.jl, LazyArrays.jl, and Thunks.jl (:top_hat: tip to Dagger.jl) - and there is also Thunks in ChainRulesCore, but not sure which approach would be best... suggestions?

view this post on Zulip Simone Carlo Surace (Jan 01 2025 at 13:51):

Do you want/need to expose the lazyness?

view this post on Zulip Simone Carlo Surace (Jan 01 2025 at 13:54):

I guess you could define your own lazy type that wraps the mlx lazy array and which can be materialized. This is how LazyArrays does it. I don‘t see why you‘d need to integrate with those Julia libraries at all.

view this post on Zulip Jesper Stemann Andersen (Jan 01 2025 at 19:29):

The input I got from Slack was to look into broadcasting, LazyArrays or how Reactant does tracing: https://julialang.slack.com/archives/C6A044SQH/p1734260542606079

view this post on Zulip Jesper Stemann Andersen (Jan 01 2025 at 19:31):

Right now an AbstractArray/StridedArray interface is implemented for MLXArray. But it seems somewhat misleading if the array is not materialized.

Last updated: Jan 29 2025 at 04:38 UTC