Skip to content

Add mlx_array_detach to break graph links before free #119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
odilitime wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add mlx_array_detach to break graph links before free #119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions mlx/c/array.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,18 @@ extern "C" int mlx_array_free(mlx_array arr) {
return 0;
}

extern "C" int mlx_array_detach(mlx_array arr) {
try {
if (arr.ctx) {
static_cast<mlx::core::array*>(arr.ctx)->detach();
}
} catch (std::exception& e) {
mlx_error(e.what());
return 1;
}
return 0;
}

extern "C" mlx_array mlx_array_new(void) {
try {
return mlx_array_();
Expand Down
10 changes: 10 additions & 0 deletions mlx/c/array.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,16 @@ mlx_array mlx_array_new(void);
*/
int mlx_array_free(mlx_array arr);

/**
* Detach an array from the MLX computation graph.
*
* Clears sibling links so the array can be freed without leaving dangling
* references in shared ArrayDesc nodes. Call before mlx_array_free when
* releasing model weights or other arrays that may still be linked in the
* graph.
*/
int mlx_array_detach(mlx_array arr);

/**
* New array from a bool scalar.
*/
Expand Down