Saving an mlx array as tiff image without having to convert to numpy first?

[kyrollosyanny]

Im running into an issue where I have an mlx model that is used for image simulation. The model is super fast ~0.1 sec per image. However, I need to save the final image as a tiff. I do this by converting the mlx ouput to a numpy array then using tiff.imwrite to write the tiff file. This increases the runtime to about 5 seconds per image. Is there a way to directly save the mlx output as an image without having to convert to numpy first?

Thank you

[angeloskath]

If I had to guess, I think that you are not evaluating the mlx array which makes it seem like the model takes 0.1 seconds per image but the computation actually happens when casting to numpy. Making a numpy copy should not take a significant amount of time. If there is memory pressure and you want to avoid copying alltogether, you can cast to numpy without copy by doing np.array(x, copy=False).

See https://ml-explore.github.io/mlx/build/html/usage/lazy_evaluation.html regarding the lazy evaluation of MLX that is probably causing the confusion I described above.

[kyrollosyanny]

Got it, Yes I wasn't running an mx.eval(). It makes sense that I'd have to do that at the end of each loop to get the true speed of my code. Thank you so much.