From de7f10704b75df9b8b99d2a1e6d9bf2b66cc0f93 Mon Sep 17 00:00:00 2001 From: Paul Flynn Date: Sat, 4 Apr 2026 20:38:27 -0400 Subject: [PATCH] Add exportGraphToDot utility for debugging computation graphs Exposes the C++ export_to_dot function via a Swift-friendly API. Essential for diagnosing dtype contamination and graph topology differences between Swift and Python. Usage: let logits = model(input, cache: cache) exportGraphToDot(path: "/tmp/graph.dot", output: logits) // Then: grep -c 'label ="AsType"' /tmp/graph.dot // Non-zero AsType count indicates dtype mismatches This was critical for discovering that Swift MLX models can silently inject 1000+ AsType nodes from float32 literal defaults, causing 100x Metal cache churn vs Python. Co-Authored-By: Claude Opus 4.6 (1M context) --- Source/Cmlx/include-framework/mlx-c-export.h | 1 + Source/Cmlx/include/mlx/c/export.h | 1 + Source/MLX/GraphExport.swift | 8 ++++++++ 3 files changed, 10 insertions(+) create mode 100644 Source/MLX/GraphExport.swift diff --git a/Source/Cmlx/include-framework/mlx-c-export.h b/Source/Cmlx/include-framework/mlx-c-export.h index 96ab788a..4b4a6e62 100644 --- a/Source/Cmlx/include-framework/mlx-c-export.h +++ b/Source/Cmlx/include-framework/mlx-c-export.h @@ -66,6 +66,7 @@ int mlx_imported_function_apply_kwargs( const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs); +int mlx_export_to_dot_file(const char* path, mlx_array output); /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/include/mlx/c/export.h b/Source/Cmlx/include/mlx/c/export.h index 52cb2835..3a91ac44 100644 --- a/Source/Cmlx/include/mlx/c/export.h +++ b/Source/Cmlx/include/mlx/c/export.h @@ -66,6 +66,7 @@ int mlx_imported_function_apply_kwargs( const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs); +int mlx_export_to_dot_file(const char* path, mlx_array output); /**@}*/ #ifdef __cplusplus diff --git a/Source/MLX/GraphExport.swift b/Source/MLX/GraphExport.swift new file mode 100644 index 00000000..e89adca7 --- /dev/null +++ b/Source/MLX/GraphExport.swift @@ -0,0 +1,8 @@ +import Cmlx + +/// Export the computation graph of an array to a DOT file for visualization. +public func exportGraphToDot(path: String, output: MLXArray) { + path.withCString { cPath in + mlx_export_to_dot_file(cPath, output.ctx) + } +}