diff --git a/x/mlxrunner/mlx/array.go b/x/mlxrunner/mlx/array.go index 675be0c86..c1871ad6c 100644 --- a/x/mlxrunner/mlx/array.go +++ b/x/mlxrunner/mlx/array.go @@ -119,7 +119,9 @@ func (t *Array) Clone() *Array { // misc. utilities -func (t *Array) Valid() bool { +// valid reports whether t still refers to an array: false once its scope +// freed it. +func (t *Array) valid() bool { return t.ctx.ctx != nil } @@ -139,7 +141,7 @@ func (t *Array) LogValue() slog.Value { attrs := []slog.Attr{ slog.String("name", t.name), } - if t.Valid() { + if t.valid() { attrs = append(attrs, slog.Any("dtype", t.DType()), slog.Any("shape", t.Dims()), diff --git a/x/mlxrunner/mlx/mlx.go b/x/mlxrunner/mlx/mlx.go index 9c34963ea..141ba9cc2 100644 --- a/x/mlxrunner/mlx/mlx.go +++ b/x/mlxrunner/mlx/mlx.go @@ -133,7 +133,7 @@ func doEval(outputs []*Array, async bool) { defer freeVectorArray(vector) for _, output := range outputs { - if output != nil && output.Valid() { + if output != nil { mlxCheck(C.mlx_vector_array_append_value(vector, output.ctx)) } } diff --git a/x/mlxrunner/mlx/nn.go b/x/mlxrunner/mlx/nn.go index d2e7fb4f1..a953f97b4 100644 --- a/x/mlxrunner/mlx/nn.go +++ b/x/mlxrunner/mlx/nn.go @@ -8,7 +8,7 @@ type Linear struct { // Forward computes the linear transformation: x @ Weight.T + Bias func (m *Linear) Forward(x *Array) *Array { w := m.Weight.Transpose(1, 0) - if m.Bias.Valid() { + if m.Bias != nil { return m.Bias.Addmm(x, w, 1.0, 1.0) } diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index d97c27a3b..125dbb528 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/x/mlxrunner/mlx/ops_extra.go @@ -153,7 +153,7 @@ func Conv1d(x, weight *Array, bias *Array, stride, padding, dilation, groups int C.int(groups), DefaultStream().ctx, )) - if bias != nil && bias.Valid() { + if bias != nil { out = Add(out, bias) } return out @@ -646,7 +646,7 @@ func collect(v reflect.Value, arrays *[]*Array, seen map[uintptr]bool) { seen[ptr] = true if arr, ok := v.Interface().(*Array); ok { - if arr != nil && arr.Valid() { + if arr != nil { *arrays = append(*arrays, arr) } return @@ -659,7 +659,7 @@ func collect(v reflect.Value, arrays *[]*Array, seen map[uintptr]bool) { case reflect.Struct: // Check if this struct IS an Array (not a pointer to one) if arr, ok := v.Addr().Interface().(*Array); ok { - if arr != nil && arr.Valid() { + if arr != nil { *arrays = append(*arrays, arr) } return diff --git a/x/mlxrunner/mlx/scope.go b/x/mlxrunner/mlx/scope.go index 92156246a..ec63d1b74 100644 --- a/x/mlxrunner/mlx/scope.go +++ b/x/mlxrunner/mlx/scope.go @@ -177,7 +177,7 @@ func (s *Scope) take(t *Array) { if from == s { return } - if !t.Valid() { + if !t.valid() { panic(fmt.Sprintf("mlx: array %q used after its scope ended", t.name)) } if from.noEscape { diff --git a/x/mlxrunner/mlx/scope_test.go b/x/mlxrunner/mlx/scope_test.go index 05994d786..641e6ddc7 100644 --- a/x/mlxrunner/mlx/scope_test.go +++ b/x/mlxrunner/mlx/scope_test.go @@ -22,17 +22,17 @@ func TestScopeFreesWhatIsNotReturned(t *testing.T) { return []*Array{FromValue(3), nil, kept} }) returned = out[0] - if !returned.Valid() { + if !returned.valid() { t.Fatal("returned array was freed with the scope that created it") } - if dropped.Valid() { + if dropped.valid() { t.Fatal("array not returned survived its scope") } }) - if returned.Valid() { + if returned.valid() { t.Fatal("returned array survived the scope it was returned into") } - if !kept.Valid() { + if !kept.valid() { t.Fatal("returning a held array moved it out of its scope") } }) @@ -48,10 +48,10 @@ func TestScopedEvalEvaluatesAfterBuild(t *testing.T) { tmp = FromValue(2) return []*Array{FromValue(1).Add(tmp)} }) - if tmp.Valid() { + if tmp.valid() { t.Fatal("intermediate survived the build scope") } - if !out[0].Valid() || out[0].Int() != 3 { + if !out[0].valid() || out[0].Int() != 3 { t.Fatal("returned array was not evaluated after the build scope") } }) @@ -72,11 +72,11 @@ func TestHeldScope(t *testing.T) { kept, discarded, detached = FromValue(1), FromValue(2), FromValue(3) held.Attach(kept, discarded, detached) held.Discard(discarded) - if discarded.Valid() { + if discarded.valid() { t.Fatal("discarded array survived") } Scoped(func() { held.Detach(detached) }) - if detached.Valid() { + if detached.valid() { t.Fatal("detached array survived the scope it was detached into") } if !panics(func() { other.Discard(kept) }) { @@ -92,11 +92,11 @@ func TestHeldScope(t *testing.T) { t.Fatal("no panic holding an array twice") } }) - if !kept.Valid() { + if !kept.valid() { t.Fatal("held array was freed with the scope that created it") } held.Close() - if kept.Valid() { + if kept.valid() { t.Fatal("held array survived its scope's close") } }) @@ -115,7 +115,7 @@ func TestScopeEndsOnPanic(t *testing.T) { panic("build failed") }) }() - if a.Valid() { + if a.valid() { t.Fatal("array survived the scope that panicked") } if currentScope != start { diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 64f5d271f..d29a20099 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -398,7 +398,7 @@ func supportsGatherQMM(mode string, bits int) bool { // weights to the [E, in, out] layout GatherMM consumes, materialized once at // load so the forward path avoids per-call transposes. func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index f9b0c6685..e402807dd 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -186,7 +186,7 @@ func sliceAxis1(a *mlx.Array, start, stop int32) *mlx.Array { // transposeForGatherMM transposes stacked expert weights from [experts, out, in] // to [experts, in, out] for use with GatherMM (which computes a @ b[group]). func transposeForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index 07751c70e..fd000e5f8 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -594,14 +594,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() } func transposeExpertWeightViewForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1) @@ -637,7 +637,7 @@ func denseExpertWeightForGatherMM(w *stackedExpertWeights) *mlx.Array { } func denseExpertWeightSupportsSourceLayout(w *stackedExpertWeights) bool { - return w != nil && w.Weight != nil && w.Weight.Valid() && w.Scales == nil && w.Weight.DType() == mlx.DTypeBFloat16 + return w != nil && w.Weight != nil && w.Scales == nil && w.Weight.DType() == mlx.DTypeBFloat16 } func denseExpertWeightsSupportSourceLayout(weights ...*stackedExpertWeights) bool { @@ -755,7 +755,7 @@ func splitLastDim(x *mlx.Array, first int32) (*mlx.Array, *mlx.Array) { } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { - if a == nil || !a.Valid() || b == nil || !b.Valid() { + if a == nil || b == nil { return nil } return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() diff --git a/x/models/nemotron_h/nemotron_h.go b/x/models/nemotron_h/nemotron_h.go index 2a8bdabf7..51584ed8b 100644 --- a/x/models/nemotron_h/nemotron_h.go +++ b/x/models/nemotron_h/nemotron_h.go @@ -422,7 +422,7 @@ func applyExpertWeightGlobalScale(weight, scale *mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() diff --git a/x/models/nn/nn.go b/x/models/nn/nn.go index 07282e911..95e5c0875 100644 --- a/x/models/nn/nn.go +++ b/x/models/nn/nn.go @@ -61,7 +61,7 @@ type Linear struct { } func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { - if bias != nil && bias.Valid() && bias.DType() != weight.DType() { + if bias != nil && bias.DType() != weight.DType() { bias = bias.AsType(weight.DType()) } return &Linear{Weight: weight, Bias: bias} @@ -69,7 +69,7 @@ func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear { func (l *Linear) Forward(x *mlx.Array) *mlx.Array { w := l.Weight.Transpose(1, 0) - if l.Bias != nil && l.Bias.Valid() { + if l.Bias != nil { return l.Bias.Addmm(x, w, 1.0, 1.0) } return x.Matmul(w) @@ -98,7 +98,7 @@ func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, } else { mlx.Eval(qw, scales) } - if bias != nil && bias.Valid() && bias.DType() != weight.DType() { + if bias != nil && bias.DType() != weight.DType() { bias = bias.AsType(weight.DType()) } return &QuantizedLinear{ @@ -130,7 +130,7 @@ func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array { // coverage for this path. out = quantizedLinearOutputScale(out, ql.GlobalScale) } - if ql.Bias != nil && ql.Bias.Valid() { + if ql.Bias != nil { bias := ql.Bias if bias.DType() != out.DType() { bias = bias.AsType(out.DType()) @@ -194,7 +194,7 @@ func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { weight := qe.Weight.TakeAxis(indices, 0) scales := qe.Scales.TakeAxis(indices, 0) var qbiases *mlx.Array - if qe.QBiases != nil && qe.QBiases.Valid() { + if qe.QBiases != nil { qbiases = qe.QBiases.TakeAxis(indices, 0) } return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode, qe.GlobalScale) diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index 64ca4dbbe..9076f8508 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -483,14 +483,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { - if w == nil || !w.Valid() || w.NumDims() != 3 { + if w == nil || w.NumDims() != 3 { return w } return mlx.Transpose(w, 0, 2, 1).Clone() } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { - if a == nil || !a.Valid() || b == nil || !b.Valid() { + if a == nil || b == nil { return nil } return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone()