Stub NAX JIT source providers when gated out

CI JIT=1 on macos-14 got past cmake configure but hit a linker
error:

  symbol not found in flat namespace '__ZN3mlx4core5metal13quantized_naxEv'

jit_kernels.cpp references mlx::core::metal::{gemm_nax,
quantized_nax, fp_quantized_nax, steel_gemm_fused_nax, …} as strong
external symbols. Those symbols are normally supplied by the
auto-generated jit/<name>.cpp files that make_jit_source() creates;
when we gate those out for old SDKs, nothing provides them and the
NIF fails to link.

Extend the patch to emit an nax_stubs.cpp with empty-string returns
under the gated-out branch. The stubs are never exercised because
is_nax_available() returns false when MLX_METAL_NO_NAX is defined.

Author: ausimian

patches/mlx-jit-nax-gate.patch

@@ -1,8 +1,8 @@
 diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt
-index 67c69579..5dfa56c1 100644
+index 67c69579..a3d575ed 100644
 --- a/mlx/backend/metal/CMakeLists.txt
 +++ b/mlx/backend/metal/CMakeLists.txt
-@@ -83,18 +83,26 @@ if(MLX_METAL_JIT)
+@@ -83,18 +83,43 @@ if(MLX_METAL_JIT)
 
    make_jit_source(steel/attn/kernels/steel_attention)
 
@@ -15,6 +15,10 @@ index 67c69579..5dfa56c1 100644
 +  # NAX JIT sources #include <MetalPerformancePrimitives/...>, which only
 +  # exists in the macOS 26.2+ SDK. Mirror the AOT gate from
 +  # kernels/CMakeLists.txt so older SDKs can still build with JIT on.
++  # When gated out we supply empty stubs for the NAX JIT source
++  # providers so jit_kernels.cpp still links; is_nax_available()
++  # returns false at runtime under MLX_METAL_NO_NAX so the stubs are
++  # never exercised.
 +  if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
 +                                                26.2))
 +    make_jit_source(
@@ -35,6 +39,19 @@ index 67c69579..5dfa56c1 100644
 +    make_jit_source(steel/attn/kernels/steel_attention_nax)
 +  else()
 +    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
++    set(_nax_stub_cpp "${CMAKE_CURRENT_BINARY_DIR}/jit/nax_stubs.cpp")
++    file(
++      WRITE "${_nax_stub_cpp}"
++      "namespace mlx::core::metal {\n"
++      "const char* gemm_nax() { return \"\"; }\n"
++      "const char* steel_gemm_fused_nax() { return \"\"; }\n"
++      "const char* steel_gemm_gather_nax() { return \"\"; }\n"
++      "const char* steel_gemm_splitk_nax() { return \"\"; }\n"
++      "const char* quantized_nax() { return \"\"; }\n"
++      "const char* fp_quantized_nax() { return \"\"; }\n"
++      "const char* steel_attention_nax() { return \"\"; }\n"
++      "}\n")
++    target_sources(mlx PRIVATE "${_nax_stub_cpp}")
 +  endif()
 
  else()