From 15f48c640510b5441c03d8761ac955b66bdf02c9 Mon Sep 17 00:00:00 2001 From: TheTom Date: Tue, 21 Apr 2026 22:08:28 -0500 Subject: [PATCH] guard mlx_array_dim against 0-dim and out-of-bounds access mlx_array_dim crashes with SmallVector out of range when called on scalar (0-dim) arrays or with dim >= ndim. This kills the process via mlx_error/fatalError during Swift metadata initialization. Returns 1 for scalars, 0 for out-of-bounds. No change for valid inputs. --- mlx/c/array.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx/c/array.cpp b/mlx/c/array.cpp index 7c7342d..8267a41 100644 --- a/mlx/c/array.cpp +++ b/mlx/c/array.cpp @@ -330,7 +330,12 @@ extern "C" const size_t* mlx_array_strides(const mlx_array arr) { } extern "C" int mlx_array_dim(const mlx_array arr, int dim) { try { - return mlx_array_get_(arr).shape(dim); + auto& a = mlx_array_get_(arr); + int ndim = a.ndim(); + if (ndim == 0) return 1; // scalar: treat as size 1 for any dim + if (dim < 0) dim += ndim; + if (dim < 0 || dim >= ndim) return 0; // out of bounds: return 0 instead of crashing + return a.shape(dim); } catch (std::exception& e) { mlx_error(e.what()); return 0;