mlxrunner: capture MLX errors in a single buffer read after every call

MLX runs on one goroutine locked to its OS thread, so the thread-local
error buffers and closure-based check helpers defended against a
calling pattern that is already invalid.

Replace them with a single buffer that the handler fills and Go reads
after every call. mlxError returns the captured message; mlxCheck
panics on it and passes the call's result through, so a checked call
is one expression. Only an int status carries a failure signal, which
lets a message next to a zero status be reported as an earlier
unchecked call.

Fix two tests that relied on errors being dropped: the laguna
mixed-precision fixture used an unsupported quantization group size,
and the compile callback test expected the callback's own panic.

x/mlxrunner/mlx/compile.go

@@ -63,9 +63,7 @@ func Compile(name string, fn CompileFunc, opts ...CompileOption) CompileFunc {
 			defer C.mlx_closure_free(src)
 
 			closure = C.mlx_closure_new()
-			mlxCheck(name+": compile failed", func() C.int {
-				return C.mlx_compile(&closure, src, C.bool(cfg.shapeless))
-			})
+			mlxCheck(C.mlx_compile(&closure, src, C.bool(cfg.shapeless)))
 		})
 
 		inVec := C.mlx_vector_array_new()
@@ -76,9 +74,7 @@ func Compile(name string, fn CompileFunc, opts ...CompileOption) CompileFunc {
 
 		outVec := C.mlx_vector_array_new()
 		defer C.mlx_vector_array_free(outVec)
-		mlxCheck(name+": closure apply failed", func() C.int {
-			return C.mlx_closure_apply(&outVec, closure, inVec)
-		})
+		mlxCheck(C.mlx_closure_apply(&outVec, closure, inVec))
 
 		n := int(C.mlx_vector_array_size(outVec))
 		outputs := make([]*Array, n)

x/mlxrunner/mlx/compile_test.go

@@ -124,8 +124,8 @@ func testCompileCallbackPanicRecovers(t *mlxthreadtest.T) {
 		if r == nil {
 			t.Fatal("expected panic from Call, got none")
 		}
-		if _, ok := r.(string); !ok {
-			t.Fatalf("expected string panic, got %T: %v", r, r)
+		if _, ok := r.(error); !ok {
+			t.Fatalf("expected error panic, got %T: %v", r, r)
 		}
 	}()
 	boom(x)

x/mlxrunner/mlx/memory.go

@@ -72,23 +72,23 @@ func ResetPeakMemory() {
 // for resident Metal allocations.
 func MaxRecommendedWorkingSetSize() (int, error) {
 	info := C.mlx_device_info_new()
-	if err := mlxCall("get device info failed", func() C.int {
-		return C.mlx_device_info_get(&info, DefaultDevice().ctx)
-	}); err != nil {
-		C.mlx_device_info_free(info)
+	defer C.mlx_device_info_free(info)
+	if err := mlxError(C.mlx_device_info_get(&info, DefaultDevice().ctx)); err != nil {
 		return 0, err
 	}
-	defer C.mlx_device_info_free(info)
 
 	key := C.CString("max_recommended_working_set_size")
 	defer C.free(unsafe.Pointer(key))
 
 	var size C.size_t
-	if err := mlxCall("max recommended working set size unavailable", func() C.int {
-		return C.mlx_device_info_get_size(&size, info, key)
-	}); err != nil {
+	rc := C.mlx_device_info_get_size(&size, info, key)
+	if err := lastError(); err != nil {
 		return 0, err
 	}
+	if rc != 0 {
+		// mlx-c reports a missing key with a non-zero return and no message.
+		return 0, fmt.Errorf("mlx: no max_recommended_working_set_size in device info")
+	}
 	return int(size), nil
 }
 
@@ -100,9 +100,7 @@ func SetWiredLimit(limit int) (int, error) {
 	}
 
 	var previous C.size_t
-	if err := mlxCall("set wired limit failed", func() C.int {
-		return C.mlx_set_wired_limit(&previous, C.size_t(limit))
-	}); err != nil {
+	if err := mlxError(C.mlx_set_wired_limit(&previous, C.size_t(limit))); err != nil {
 		return 0, err
 	}
 	return int(previous), nil

x/mlxrunner/mlx/mlx.go

@@ -1,3 +1,8 @@
+// Package mlx wraps the MLX C API.
+//
+// MLX keeps stream and backend state in thread-locals, so all calls into this
+// package must come from a single goroutine locked to its OS thread (see
+// x/internal/mlxthread).
 package mlx
 
 //go:generate go run generator/main.go -output=. ./include/mlx/c/*.h
@@ -9,35 +14,27 @@ package mlx
 // #include "generated.h"
 // #include <string.h>
 //
-// static __thread char _mlx_last_error_msg[1024] = {0};
-// static __thread int  _mlx_last_error_flag = 0;
+// static char _mlx_last_error[1024];
 //
-// static void _mlx_capture_error_handler(const char* msg, void* data) {
+// static void _mlx_capture_error(const char* msg, void* data) {
 //     (void)data;
-//     strncpy(_mlx_last_error_msg, msg, sizeof(_mlx_last_error_msg) - 1);
-//     _mlx_last_error_msg[sizeof(_mlx_last_error_msg) - 1] = '\0';
-//     _mlx_last_error_flag = 1;
+//     strncpy(_mlx_last_error, msg, sizeof(_mlx_last_error) - 1);
 // }
 //
 // static void mlx_install_capture_handler(void) {
 //     if (mlx_set_error_handler_) {
-//         mlx_set_error_handler_(_mlx_capture_error_handler, NULL, NULL);
+//         mlx_set_error_handler_(_mlx_capture_error, NULL, NULL);
 //     }
 // }
 //
-// static void mlx_clear_last_error(void) {
-//     _mlx_last_error_flag = 0;
-//     _mlx_last_error_msg[0] = '\0';
-// }
-//
-// static const char* mlx_get_last_error(void) {
-//     return _mlx_last_error_flag ? _mlx_last_error_msg : "";
+// static char* mlx_last_error(void) {
+//     return _mlx_last_error;
 // }
 import "C"
 
 import (
+	"errors"
 	"fmt"
-	"runtime"
 )
 
 func init() {
@@ -46,6 +43,56 @@ func init() {
 	C.mlx_install_capture_handler()
 }
 
+var errBuf = C.mlx_last_error()
+
+// lastError consumes the captured MLX error, or returns nil when none is
+// pending.
+func lastError() error {
+	if *errBuf == 0 {
+		return nil
+	}
+	err := fmt.Errorf("mlx: %s", C.GoString(errBuf))
+	*errBuf = 0
+	return err
+}
+
+// mlxError returns the MLX error captured by the call that produced v. mlx-c
+// signals failure with a non-zero int status; a message next to a zero
+// status came from an earlier unchecked call.
+func mlxError[T comparable](v T) error {
+	var zero T
+	var failed, signaled bool
+	switch any(zero).(type) {
+	case C.int:
+		failed, signaled = v != zero, true
+	default:
+		// Only an int status signals failure. Handles, pointers, sizes, and
+		// dtypes are all valid at zero: a null handle is what the out-param
+		// constructors return, and an empty array has no data.
+	}
+	if *errBuf != 0 {
+		err := lastError()
+		if signaled && !failed {
+			return fmt.Errorf("mlx: unchecked error from an earlier call: %w", err)
+		}
+		return err
+	}
+	if failed {
+		return errors.New("mlx: call failed without an error message")
+	}
+	return nil
+}
+
+// mlxCheck panics on a failed call and otherwise passes its result through.
+// Most array operations cannot recover from a failed graph construction or
+// evaluation.
+func mlxCheck[T comparable](v T) T {
+	if err := mlxError(v); err != nil {
+		panic(err)
+	}
+	return v
+}
+
 // Version returns the MLX core library version string.
 func Version() string {
 	str := C.mlx_string_new()
@@ -54,31 +101,6 @@ func Version() string {
 	return C.GoString(C.mlx_string_data(str))
 }
 
-// mlxCall locks the goroutine to its OS thread so the thread-local error state
-// is read from the same thread that executed fn.
-func mlxCall(fallback string, fn func() C.int) error {
-	runtime.LockOSThread()
-	defer runtime.UnlockOSThread()
-
-	C.mlx_clear_last_error()
-	if fn() != 0 {
-		msg := C.GoString(C.mlx_get_last_error())
-		if msg == "" {
-			msg = fallback
-		}
-		return fmt.Errorf("mlx: %s", msg)
-	}
-	return nil
-}
-
-// mlxCheck panics with the captured MLX error. Most array operations cannot
-// recover from a failed graph construction or evaluation.
-func mlxCheck(fallback string, fn func() C.int) {
-	if err := mlxCall(fallback, fn); err != nil {
-		panic(err.Error())
-	}
-}
-
 func doEval(outputs []*Array, async bool) {
 	if len(outputs) == 0 {
 		return
@@ -93,12 +115,11 @@ func doEval(outputs []*Array, async bool) {
 		}
 	}
 
-	mlxCheck("eval failed", func() C.int {
-		if async {
-			return C.mlx_async_eval(vector)
-		}
-		return C.mlx_eval(vector)
-	})
+	if async {
+		mlxCheck(C.mlx_async_eval(vector))
+	} else {
+		mlxCheck(C.mlx_eval(vector))
+	}
 }
 
 func AsyncEval(outputs ...*Array) {

x/models/laguna/laguna_test.go

@@ -197,7 +197,7 @@ func TestTinyLagunaLoadAndForward(t *testing.T) {
 	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
-		"hidden_size": 8,
+		"hidden_size": 32,
 		"intermediate_size": 12,
 		"moe_intermediate_size": 4,
 		"shared_expert_intermediate_size": 4,
@@ -264,8 +264,8 @@ func TestTinyLagunaLoadAndForward(t *testing.T) {
 			SeqQueryLens: []int32{int32(tokens.Dim(1))},
 		}, caches)
 		mlx.Eval(hidden)
-		if got := hidden.Dims(); len(got) != 3 || got[0] != 1 || got[1] != 3 || got[2] != 8 {
-			t.Fatalf("hidden shape = %v, want [1 3 8]", got)
+		if got := hidden.Dims(); len(got) != 3 || got[0] != 1 || got[1] != 3 || got[2] != 32 {
+			t.Fatalf("hidden shape = %v, want [1 3 32]", got)
 		}
 
 		logits := m.Unembed(hidden)
@@ -285,7 +285,7 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
 	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
-		"hidden_size": 8,
+		"hidden_size": 32,
 		"intermediate_size": 12,
 		"moe_intermediate_size": 4,
 		"shared_expert_intermediate_size": 4,
@@ -329,7 +329,7 @@ func TestTinyLagunaLoadWeightsFusesDenseGateUp(t *testing.T) {
 		if moe.SwitchMLP.GateUpWeight == nil {
 			t.Fatal("expected fused GateUpWeight to be populated")
 		}
-		if got, want := moe.SwitchMLP.GateUpWeight.Dims(), []int{2, 8, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
+		if got, want := moe.SwitchMLP.GateUpWeight.Dims(), []int{2, 32, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
 			t.Fatalf("GateUpWeight dims = %v, want %v", got, want)
 		}
 	})
@@ -339,7 +339,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
 	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg, err := parseConfig([]byte(`{
 		"model_type": "laguna",
-		"hidden_size": 8,
+		"hidden_size": 32,
 		"intermediate_size": 12,
 		"moe_intermediate_size": 4,
 		"shared_expert_intermediate_size": 4,
@@ -392,7 +392,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
 		if moe.SwitchMLP.GateUpWeight != nil {
 			t.Fatal("expected BF16 source-layout SwitchMLP to avoid pre-fused gate/up weights")
 		}
-		if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 8}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
+		if got, want := moe.SwitchMLP.GateWeight.Dims(), []int{2, 4, 32}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] || got[2] != want[2] {
 			t.Fatalf("GateWeight dims = %v, want %v", got, want)
 		}
 	})
@@ -401,7 +401,7 @@ func TestTinyLagunaLoadWeightsKeepsBF16SourceLayout(t *testing.T) {
 func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) {
 	mlxtest.Run(t, func(t *mlxtest.T) {
 		cfg := &Config{
-			HiddenSize:                8,
+			HiddenSize:                32,
 			IntermediateSize:          12,
 			MoeIntermediateSize:       4,
 			SharedExpertIntermediate:  4,
@@ -418,7 +418,7 @@ func TestTinyLagunaLoadWeightsKeepsMixedExpertPrecision(t *testing.T) {
 			NumExpertsPerTok:          1,
 			MoeRoutedScalingFactor:    2.5,
 			RMSNormEps:                1e-5,
-			QuantGroupSize:            4,
+			QuantGroupSize:            32,
 			QuantBits:                 4,
 			QuantMode:                 "affine",
 		}
@@ -739,38 +739,38 @@ func TestCombinedTensorGlobalScaleIgnoresInputGlobalScale(t *testing.T) {
 
 func tinyLagunaTensors() map[string]*mlx.Array {
 	tensors := map[string]*mlx.Array{
-		"model.embed_tokens.weight": weights(16, 8),
-		"model.norm.weight":         ones(8),
-		"lm_head.weight":            weights(16, 8),
+		"model.embed_tokens.weight": weights(16, 32),
+		"model.norm.weight":         ones(32),
+		"lm_head.weight":            weights(16, 32),
 	}
 	for layer := range 2 {
 		prefix := "model.layers." + string(rune('0'+layer))
-		tensors[prefix+".input_layernorm.weight"] = ones(8)
-		tensors[prefix+".post_attention_layernorm.weight"] = ones(8)
-		tensors[prefix+".self_attn.q_proj.weight"] = weights(8, 8)
-		tensors[prefix+".self_attn.k_proj.weight"] = weights(4, 8)
-		tensors[prefix+".self_attn.v_proj.weight"] = weights(4, 8)
-		tensors[prefix+".self_attn.o_proj.weight"] = weights(8, 8)
-		tensors[prefix+".self_attn.g_proj.weight"] = weights(2, 8)
+		tensors[prefix+".input_layernorm.weight"] = ones(32)
+		tensors[prefix+".post_attention_layernorm.weight"] = ones(32)
+		tensors[prefix+".self_attn.q_proj.weight"] = weights(8, 32)
+		tensors[prefix+".self_attn.k_proj.weight"] = weights(4, 32)
+		tensors[prefix+".self_attn.v_proj.weight"] = weights(4, 32)
+		tensors[prefix+".self_attn.o_proj.weight"] = weights(32, 8)
+		tensors[prefix+".self_attn.g_proj.weight"] = weights(2, 32)
 		tensors[prefix+".self_attn.q_norm.weight"] = ones(4)
 		tensors[prefix+".self_attn.k_norm.weight"] = ones(4)
 	}
 
-	tensors["model.layers.0.mlp.gate_proj.weight"] = weights(12, 8)
-	tensors["model.layers.0.mlp.up_proj.weight"] = weights(12, 8)
-	tensors["model.layers.0.mlp.down_proj.weight"] = weights(8, 12)
+	tensors["model.layers.0.mlp.gate_proj.weight"] = weights(12, 32)
+	tensors["model.layers.0.mlp.up_proj.weight"] = weights(12, 32)
+	tensors["model.layers.0.mlp.down_proj.weight"] = weights(32, 12)
 
-	tensors["model.layers.1.mlp.gate.weight"] = weights(2, 8)
+	tensors["model.layers.1.mlp.gate.weight"] = weights(2, 32)
 	tensors["model.layers.1.mlp.experts.e_score_correction_bias"] = mlx.FromValues([]float32{0.1, -0.1}, 2)
 	for expert := range 2 {
 		prefix := "model.layers.1.mlp.experts." + string(rune('0'+expert))
-		tensors[prefix+".gate_proj.weight"] = weights(4, 8)
-		tensors[prefix+".up_proj.weight"] = weights(4, 8)
-		tensors[prefix+".down_proj.weight"] = weights(8, 4)
+		tensors[prefix+".gate_proj.weight"] = weights(4, 32)
+		tensors[prefix+".up_proj.weight"] = weights(4, 32)
+		tensors[prefix+".down_proj.weight"] = weights(32, 4)
 	}
-	tensors["model.layers.1.mlp.shared_expert.gate_proj.weight"] = weights(4, 8)
-	tensors["model.layers.1.mlp.shared_expert.up_proj.weight"] = weights(4, 8)
-	tensors["model.layers.1.mlp.shared_expert.down_proj.weight"] = weights(8, 4)
+	tensors["model.layers.1.mlp.shared_expert.gate_proj.weight"] = weights(4, 32)
+	tensors["model.layers.1.mlp.shared_expert.up_proj.weight"] = weights(4, 32)
+	tensors["model.layers.1.mlp.shared_expert.down_proj.weight"] = weights(32, 4)
 	return tensors
 }