Saving and Loading Arrays

Learn how to serialize and deserialize MLX arrays using NumPy, Safetensors, and GGUF formats

MLX supports multiple array serialization formats for saving and loading arrays.

Serialization Formats

Format	Extension	Function	Notes
NumPy	.npy	mx.save()	Single arrays only
NumPy archive	.npz	mx.savez() and mx.savez_compressed()	Multiple arrays
Safetensors	.safetensors	mx.save_safetensors()	Multiple arrays
GGUF	.gguf	mx.save_gguf()	Multiple arrays

The mx.load() function will load any of the supported serialization formats. It determines the format from the extensions. The output of mx.load() depends on the format.

Saving Single Arrays

Here’s an example of saving a single array to a file:

a = mx.array([1.0])
mx.save("array", a)

The array a will be saved in the file array.npy (notice the extension is automatically added). Including the extension is optional; if it is missing it will be added.

1. Save the array
2. a = mx.array([1.0])
   mx.save("array", a)
   

3. Load the array
4. mx.load("array.npy")
   # array([1], dtype=float32)
   

Saving Multiple Arrays

Here’s an example of saving several arrays to a single file:

a = mx.array([1.0])
b = mx.array([2.0])
mx.savez("arrays", a, b=b)

For compatibility with numpy.savez, the MLX savez() takes arrays as arguments. If the keywords are missing, then default names will be provided.

Loading Multiple Arrays

This can be loaded with:

mx.load("arrays.npz")
# {'b': array([2], dtype=float32), 'arr_0': array([1], dtype=float32)}

In this case mx.load() returns a dictionary of names to arrays.

Safetensors and GGUF

The functions mx.save_safetensors() and mx.save_gguf() are similar to mx.savez(), but they take as input a dict of string names to arrays:

a = mx.array([1.0])
b = mx.array([2.0])
mx.save_safetensors("arrays", {"a": a, "b": b})

Use Safetensors format for better compatibility with Hugging Face models and for improved safety guarantees.