Recommended method for control flow with mlx types

[kyrollosyanny]

Given that mlx uses lazy evaluation, using mlx arrays for control flow can be very inefficient. I was wondering if there is a recommended way for doing something as below:

a=mx.array(5)
if a==0:
    s=mx.array(5)
else:
    s= 1/a 

Thanks a lot for the great package.

[angeloskath]

The way would be using mx.where as follows

a = mx.array(5)
s = mx.where(a == 0, mx.array(5), 1/a)

[kyrollosyanny]

wow, that's much faster now! Thanks for the super fast response!