mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-25 18:11:15 +08:00
82 lines
2.1 KiB
ReStructuredText
82 lines
2.1 KiB
ReStructuredText
.. _saving_and_loading:
|
|
|
|
Saving and Loading Arrays
|
|
=========================
|
|
|
|
.. currentmodule:: mlx.core
|
|
|
|
MLX supports multiple array serialization formats.
|
|
|
|
.. list-table:: Serialization Formats
|
|
:widths: 20 8 25 25
|
|
:header-rows: 1
|
|
|
|
* - Format
|
|
- Extension
|
|
- Function
|
|
- Notes
|
|
* - NumPy
|
|
- ``.npy``
|
|
- :func:`save`
|
|
- Single arrays only
|
|
* - NumPy archive
|
|
- ``.npz``
|
|
- :func:`savez` and :func:`savez_compressed`
|
|
- Multiple arrays
|
|
* - Safetensors
|
|
- ``.safetensors``
|
|
- :func:`save_safetensors`
|
|
- Multiple arrays
|
|
* - GGUF
|
|
- ``.gguf``
|
|
- :func:`save_gguf`
|
|
- Multiple arrays
|
|
|
|
The :func:`load` function will load any of the supported serialization
|
|
formats. It determines the format from the extensions. The output of
|
|
:func:`load` depends on the format.
|
|
|
|
Here's an example of saving a single array to a file:
|
|
|
|
.. code-block:: shell
|
|
|
|
>>> 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. You can load the array with:
|
|
|
|
.. code-block:: shell
|
|
|
|
>>> mx.load("array.npy")
|
|
array([1], dtype=float32)
|
|
|
|
Here's an example of saving several arrays to a single file:
|
|
|
|
.. code-block:: shell
|
|
|
|
>>> a = mx.array([1.0])
|
|
>>> b = mx.array([2.0])
|
|
>>> mx.savez("arrays", a, b=b)
|
|
|
|
For compatibility with :func:`numpy.savez` the MLX :func:`savez` takes arrays
|
|
as arguments. If the keywords are missing, then default names will be
|
|
provided. This can be loaded with:
|
|
|
|
.. code-block:: shell
|
|
|
|
>>> mx.load("arrays.npz")
|
|
{'b': array([2], dtype=float32), 'arr_0': array([1], dtype=float32)}
|
|
|
|
In this case :func:`load` returns a dictionary of names to arrays.
|
|
|
|
The functions :func:`save_safetensors` and :func:`save_gguf` are similar to
|
|
:func:`savez`, but they take as input a :obj:`dict` of string names to arrays:
|
|
|
|
.. code-block:: shell
|
|
|
|
>>> a = mx.array([1.0])
|
|
>>> b = mx.array([2.0])
|
|
>>> mx.save_safetensors("arrays", {"a": a, "b": b})
|