Skip to content

add saveToData and load from data (safetensors) -- fix #214#218

Merged
davidkoski merged 4 commits intomlx0_24_2from
safetensor_data
May 1, 2025
Merged

add saveToData and load from data (safetensors) -- fix #214#218
davidkoski merged 4 commits intomlx0_24_2from
safetensor_data

Conversation

@davidkoski
Copy link
Collaborator

This saves and loads from memory:

        let arrays: [String: MLXArray] = [
            "foo": MLX.ones([1, 2]),
            "bar": MLX.zeros([2, 1]),
        ]

        let data = try saveToData(arrays: arrays)
        let loadedArrays = try loadArrays(data: data)

@davidkoski davidkoski requested review from awni and barronalex April 7, 2025 18:23
import Cmlx
import Foundation

/// Save array to a binary file in `.npy`format.
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This first chunk is just moving related methods from Ops.swift

}
}

// MARK: - Memory I/O
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/// Save dictionary of arrays in `safetensors` format into `Data`.
///
/// - Parameters:
/// - a: array to save
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: incorrectly labeled should be arrays

/// - Parameters:
/// - a: array to save
/// - metadata: metadata to save
/// - stream: stream or device to evaluate on
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Am I missing something or is the stream unused?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are correct -- the load produces something that goes on a stream, not the save

Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

@awni
Copy link
Member

awni commented Apr 28, 2025

Just left a couple comments, please take a look and merge when ready.

@davidkoski davidkoski merged commit 0e7cd4c into mlx0_24_2 May 1, 2025
1 check was pending
@davidkoski davidkoski deleted the safetensor_data branch May 1, 2025 15:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants