From 9b4b828825c1a1bd6dbb72c060da6cf585126fbc Mon Sep 17 00:00:00 2001 From: Lucius Carr Date: Mon, 26 Oct 2026 00:34:28 -0700 Subject: [PATCH] [Metal] Fix fill loop in `layer_norm` Co-authored-by: Arup Kanti Dash <30518803+ArupKantiDas@users.noreply.github.com> --- mlx/backend/metal/kernels/layer_norm.metal | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/kernels/layer_norm.metal b/mlx/backend/metal/kernels/layer_norm.metal index ea77b53dca..0f00f5506d 100644 --- a/mlx/backend/metal/kernels/layer_norm.metal +++ b/mlx/backend/metal/kernels/layer_norm.metal @@ -73,7 +73,7 @@ template // Compute some variables for reading writing etc const bool safe = lid * N_READS + N_READS <= axis_size; - const int n = axis_size - lid * N_READS; + const int n = max(int(axis_size) - int(lid * N_READS), 0); // Read the inputs if (safe) { @@ -235,7 +235,7 @@ template // Compute some variables for reading writing etc const bool safe = lid * N_READS + N_READS <= axis_size; - const int n = axis_size - lid * N_READS; + const int n = max(int(axis_size) - int(lid * N_READS), 0); // Read the inputs if (safe) {