zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Export / import functions to / from a file (#1642)

Browse Source 
* export and import functions

* refactor + works for few primitives

* nit

* allow primitives with state

* nit

* nit

* simplify serialize / deserialize

* fix for constants

* python bindings

* maybe fix serialize failure case

* add example

* more primitives, training kind of works

* same result for python and c++

* some fixes

* fix export

* template it up

* some simplificatoin

* rebase

* allow kwargs and multiple functions

* exporter

* more primitives for exporting

* deal with endianness

* handle invalid stream

* add docstring
This commit is contained in:
Awni Hannun
2024-12-24 11:19:13 -08:00
committed by GitHub
parent 935c8c4bb1
commit 4ba0c24a8f
35 changed files with 2239 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
mlx/io/load.h
View File

@@ -78,8 +78,15 @@ class ParallelFileReader : public Reader {
return lseek(fd_, 0, SEEK_CUR);
}
void seek(int64_t, std::ios_base::seekdir = std::ios_base::beg) override {
throw std::runtime_error("[ParallelFileReader::seek] Not allowed");
// Warning: do not use this function from multiple threads as
// it advances the file descriptor
void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
override {
if (way == std::ios_base::beg) {
lseek(fd_, off, 0);
} else {
lseek(fd_, off, SEEK_CUR);
}
}
// Warning: do not use this function from multiple threads as
@@ -108,8 +115,16 @@ class FileWriter : public Writer {
0644)),
label_(std::move(file_path)) {}
FileWriter(const FileWriter&) = delete;
FileWriter& operator=(const FileWriter&) = delete;
FileWriter(FileWriter&& other) {
std::swap(fd_, other.fd_);
}
~FileWriter() override {
close(fd_);
if (fd_ != 0) {
close(fd_);
}
}
bool is_open() const override {
@@ -151,7 +166,7 @@ class FileWriter : public Writer {
}
private:
int fd_;
int fd_{0};
std::string label_;
};