Skip to content

Debug time graph verification - #4711

Open
zcbenz wants to merge 1 commit into
ml-explore:mainfrom
zcbenz:graph-leak-check
Open

zcbenz wants to merge 1 commit into
ml-explore:mainfrom
zcbenz:graph-leak-check

Conversation

@zcbenz

@zcbenz zcbenz commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

Add debug time check that would throws error when 2 common types of errors are found:

1. Temporary arrays are not registered by add_temporary

This is checked by iterating all pointers in the graph and checking whether they are valid allocations, which can capture mistakes like the ones fixed by #4660:

Traceback (most recent call last):
  File "/Users/zcbenz/codes/mlx/python/tests/test_fast_gated_delta.py", line 317, in run
    mx.eval(outs, vjps)
    ~~~~~~~^^^^^^^^^^^^
RuntimeError: There are invalid memory registered when calling GatedDeltaUpdateVJP::eval_gpu, which is a bug that corrupts the graph, and usually caused by temporary arrays not registered with "add_temporray" method.

2. Input or outputs arrays are accidentally registered by add_temporary

This is checked by simply checking whether the temporaries contain inputs or outputs of a primitive, which can capture mistakes like the ones fixed by #4702:

Traceback (most recent call last):
  File "/Users/zcbenz/codes/mlx/issue.py", line 17, in <module>
    mx.eval(mx.array_equal(g, expected))
    ~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: A output of array is registered as temporary when calling CrossEntropyVJP::eval_gpu, which is a bug that corrupts the graph.

@zcbenz
zcbenz force-pushed the graph-leak-check branch 2 times, most recently from 36efc70 to 99b0958 Compare October 17, 2026 05:07

This branch has not been deployed

No deployments
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.

1 participant