feat: initial release

Assisted-by: GLM 5.3 Flash

grad/grad_guard_pins_test.go

@@ -0,0 +1,875 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package grad
+
+import (
+	"math"
+	"strings"
+	"testing"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+// Regression pins for the grad package guards: refusals and gradients
+// that used to panic or silently pass through. Each test names
+// the defect it pins and fails without its fix.
+
+// TestConcatMixedDtypeBackwardNarrows pins the mixed real/complex
+// Concat backward. The join promotes along the dtype ladder, so a
+// complex gradient reaches a real operand; that operand must receive
+// 2·Re of its span, the package's real-operand rule, instead of dying
+// in copyElem, which used to read the complex source through FloatAt (a
+// nil int payload) and panic. Both operand orders and the constant
+// complex operand are covered, and the real side is checked against
+// central differences as well as its closed form.
+func TestConcatMixedDtypeBackwardNarrows(t *testing.T) {
+	xv := []float64{1.5, -0.75, 2.25, 0.5, -1, 3}
+	zv := []complex128{1 + 1i, 2 - 0.5i, -0.25 + 0.75i, 3, 0.5 - 2i, -1.5i}
+
+	// L = Σ|Concat(a, b)|² = Σx² + Σ|z|², so the real operand's
+	// gradient is 2x and the complex one's is z, whatever the axis.
+	cases := []struct {
+		name string
+		dim  int
+		xSh  []int
+		zSh  []int
+	}{
+		{name: "dim0", dim: 0, xSh: []int{2, 3}, zSh: []int{2, 3}},
+		{name: "dim1", dim: 1, xSh: []int{3, 2}, zSh: []int{3, 2}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			x, err := core.FromFloats(xv, tc.xSh...)
+			if err != nil {
+				t.Fatalf("FromFloats: %v", err)
+			}
+			z, err := core.FromComplexes(zv, tc.zSh...)
+			if err != nil {
+				t.Fatalf("FromComplexes: %v", err)
+			}
+
+			// The A side of a real x complex join is the panic the
+			// report reproduced; the B side of a complex x real join
+			// fails identically.
+			concatChecked(t, "real||cplx", FromArray(x, true), FromArray(z, true), tc.dim, xv, zv)
+			concatChecked(t, "cplx||real", FromArray(z, true), FromArray(x, true), tc.dim, xv, zv)
+
+			// A constant complex operand takes the same span loop, so
+			// the real side must still narrow.
+			concatChecked(t, "real||cplxconst", FromArray(x, true), FromArray(z, false), tc.dim, xv, nil)
+
+			// Finite differences confirm the 2·Re rule itself, not just
+			// its agreement with the closed form.
+			ref := numericGrad(func(v *core.Array) float64 {
+				realSide := FromArray(v, false)
+				cat, cerr := realSide.Concat(FromArray(z, false), tc.dim)
+				if cerr != nil {
+					return math.NaN()
+				}
+				sq, cerr := cat.Abs2()
+				if cerr != nil {
+					return math.NaN()
+				}
+				s, cerr := sq.Sum()
+				if cerr != nil {
+					return math.NaN()
+				}
+				return s.Data().FloatAt(0)
+			}, x)
+			xt := FromArray(x, true)
+			concatChecked(t, "fd", xt, FromArray(z, false), tc.dim, xv, nil)
+			if d := maxAbsDiff(flatFloats(xt.Grad()), ref); d > 1e-8 {
+				t.Errorf("real operand gradient differs from central differences by %g", d)
+			}
+		})
+	}
+
+	// A float32 real operand beside a complex one: the narrowed side
+	// keeps the operand's width, and the value is still 2x.
+	t.Run("float32RealSide", func(t *testing.T) {
+		x32 := []float32{1.5, -0.75, 2.25, 0.5}
+		z32 := []complex128{1 + 1i, 2 - 0.5i, -0.25 + 0.75i, 3}
+		a, err := core.FromFloat32s(x32, 2, 2)
+		if err != nil {
+			t.Fatalf("FromFloat32s: %v", err)
+		}
+		z, err := core.FromComplexes(z32, 2, 2)
+		if err != nil {
+			t.Fatalf("FromComplexes: %v", err)
+		}
+		xt := FromArray(a, true)
+		cat, err := xt.Concat(FromArray(z, true), 0)
+		if err != nil {
+			t.Fatalf("Concat: %v", err)
+		}
+		sq, err := cat.Abs2()
+		if err != nil {
+			t.Fatalf("Abs2: %v", err)
+		}
+		loss, err := sq.Sum()
+		if err != nil {
+			t.Fatalf("Sum: %v", err)
+		}
+		if err := loss.Backward(); err != nil {
+			t.Fatalf("Backward: %v", err)
+		}
+		g := xt.Grad()
+		if got := g.Dtype(); got != core.Float32 {
+			t.Errorf("float32 real operand gradient dtype = %s, want float32", got)
+		}
+		for i, v := range x32 {
+			if got, want := g.FloatAt(i), 2*float64(v); math.Abs(got-want) > 1e-6 {
+				t.Errorf("float32 gradient[%d] = %v, want %v (2·Re rule)", i, got, want)
+			}
+		}
+	})
+
+	// A rebased-view complex operand: the view aliases a longer payload,
+	// and the gradient must scatter back through the view's own slots.
+	t.Run("rebasedViewOperand", func(t *testing.T) {
+		big, err := core.FromComplexes(zv[:4], 4)
+		if err != nil {
+			t.Fatalf("FromComplexes: %v", err)
+		}
+		zb := FromArray(big, true)
+		view, err := zb.Slice(0, 1, 3)
+		if err != nil {
+			t.Fatalf("Slice: %v", err)
+		}
+		x, err := core.FromFloats(xv[:2], 2)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		xt := FromArray(x, true)
+		cat, err := xt.Concat(view, 0)
+		if err != nil {
+			t.Fatalf("Concat: %v", err)
+		}
+		sq, err := cat.Abs2()
+		if err != nil {
+			t.Fatalf("Abs2: %v", err)
+		}
+		loss, err := sq.Sum()
+		if err != nil {
+			t.Fatalf("Sum: %v", err)
+		}
+		if err := loss.Backward(); err != nil {
+			t.Fatalf("Backward: %v", err)
+		}
+		for i := range xv[:2] {
+			if got, want := xt.Grad().FloatAt(i), 2*xv[i]; got != want {
+				t.Errorf("view case: real gradient[%d] = %v, want %v", i, got, want)
+			}
+		}
+		g := zb.Grad()
+		if g == nil {
+			t.Fatal("the view's parent received no gradient")
+		}
+		// Only the aliased slots carry the view's own values; the rest
+		// stays zero through the Slice backward.
+		want := []complex128{0, zv[1], zv[2], 0}
+		for i := range want {
+			if got := g.ComplexAt(i); got != want[i] {
+				t.Errorf("view case: parent gradient[%d] = %v, want %v", i, got, want[i])
+			}
+		}
+	})
+
+	// The mirror image: a rebased-view REAL operand beside a complex one.
+	// The narrowed span (2·Re) reaches the view as float and Slice's
+	// backward scatters it into the parent's own slots, leaving the rest
+	// zero. This is the one shape that exercises both fixes together.
+	t.Run("rebasedViewRealOperand", func(t *testing.T) {
+		full := []float64{9, 1.5, -0.75, 7}
+		fa, err := core.FromFloats(full, 4)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		fb := FromArray(fa, true)
+		view, err := fb.Slice(0, 1, 3)
+		if err != nil {
+			t.Fatalf("Slice: %v", err)
+		}
+		z, err := core.FromComplexes(zv[:2], 2)
+		if err != nil {
+			t.Fatalf("FromComplexes: %v", err)
+		}
+		cat, err := view.Concat(FromArray(z, true), 0)
+		if err != nil {
+			t.Fatalf("Concat: %v", err)
+		}
+		sq, err := cat.Abs2()
+		if err != nil {
+			t.Fatalf("Abs2: %v", err)
+		}
+		loss, err := sq.Sum()
+		if err != nil {
+			t.Fatalf("Sum: %v", err)
+		}
+		if err := loss.Backward(); err != nil {
+			t.Fatalf("Backward: %v", err)
+		}
+		g := fb.Grad()
+		if g == nil {
+			t.Fatal("the view's parent received no gradient")
+		}
+		want := []float64{0, 2 * full[1], 2 * full[2], 0}
+		for i := range want {
+			if got := g.FloatAt(i); got != want[i] {
+				t.Errorf("view case: parent gradient[%d] = %v, want %v", i, got, want[i])
+			}
+		}
+	})
+}
+
+// concatChecked builds Σ|first.Concat(second, dim)|², runs the
+// backward and checks every gradient-carrying operand against its
+// closed form by dtype: a real side against 2x (the 2·Re rule), a
+// complex side against z. A nil wantZ marks a constant complex operand,
+// and a non-grad operand has nothing to check.
+func concatChecked(t *testing.T, label string, first, second *Tensor, dim int, xv []float64, wantZ []complex128) {
+	t.Helper()
+	cat, err := first.Concat(second, dim)
+	if err != nil {
+		t.Fatalf("%s: Concat: %v", label, err)
+	}
+	sq, err := cat.Abs2()
+	if err != nil {
+		t.Fatalf("%s: Abs2: %v", label, err)
+	}
+	loss, err := sq.Sum()
+	if err != nil {
+		t.Fatalf("%s: Sum: %v", label, err)
+	}
+	if err := loss.Backward(); err != nil {
+		t.Fatalf("%s: Backward: %v", label, err)
+	}
+	for _, side := range []*Tensor{first, second} {
+		if !side.RequiresGrad() {
+			continue
+		}
+		g := side.Grad()
+		if g == nil {
+			t.Fatalf("%s: an operand received no gradient", label)
+		}
+		if g.Dtype() == core.Complex {
+			if wantZ == nil {
+				continue
+			}
+			for i := range wantZ {
+				if got := g.ComplexAt(i); got != wantZ[i] {
+					t.Errorf("%s: complex operand gradient[%d] = %v, want %v", label, i, got, wantZ[i])
+				}
+			}
+			continue
+		}
+		for i := range xv {
+			if got, want := g.FloatAt(i), 2*xv[i]; got != want {
+				t.Errorf("%s: real operand gradient[%d] = %v, want %v (the 2·Re rule)", label, i, got, want)
+			}
+		}
+	}
+}
+
+// TestHessianAndHVPRejectComplexOperands pins the dtype guard on the
+// two second-order helpers: a complex point or direction is refused
+// with an error naming the dtype, as MinimiseNewtonCG, SampleHMC and
+// AdjointODE already do, never a panic out of flatFloats.
+func TestHessianAndHVPRejectComplexOperands(t *testing.T) {
+	z, err := core.FromComplexes([]complex128{1 + 1i, 2 - 1i}, 2)
+	if err != nil {
+		t.Fatalf("FromComplexes: %v", err)
+	}
+	zt := FromArray(z, false)
+	xf, err := core.FromFloats([]float64{1, 2}, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	xt := FromArray(xf, false)
+	objective := func(q *Tensor) (*Tensor, error) {
+		sq, err := q.Abs2()
+		if err != nil {
+			return nil, err
+		}
+		return sq.Sum()
+	}
+
+	refusesComplex(t, "Hessian on a complex point", func() error {
+		_, err := Hessian(objective, zt, HessianOptions{})
+		return err
+	})
+	refusesComplex(t, "HessianVectorProduct on a complex point", func() error {
+		_, err := HessianVectorProduct(objective, zt, xt, HessianOptions{})
+		return err
+	})
+	refusesComplex(t, "HessianVectorProduct on a complex direction", func() error {
+		_, err := HessianVectorProduct(objective, xt, zt, HessianOptions{})
+		return err
+	})
+
+	// The guards must not narrow the accepted surface: a real point
+	// still differentiates, and the real callers keep working.
+	h, err := Hessian(objective, xt, HessianOptions{})
+	if err != nil {
+		t.Fatalf("Hessian of a real point: %v", err)
+	}
+	// Σ|q|² over a real q is Σq², whose Hessian is 2·I.
+	for i := range 2 {
+		if got := h.FloatAt(i*2 + i); math.Abs(got-2) > 1e-6 {
+			t.Errorf("real Hessian diagonal[%d] = %v, want 2", i, got)
+		}
+	}
+}
+
+// refusesComplex runs fn and requires an error that names the
+// complex dtype, treating a panic as the failure it reports.
+func refusesComplex(t *testing.T, label string, fn func() error) {
+	t.Helper()
+	defer func() {
+		if r := recover(); r != nil {
+			t.Errorf("%s panicked: %v", label, r)
+		}
+	}()
+	err := fn()
+	if err == nil {
+		t.Errorf("%s was accepted", label)
+		return
+	}
+	if !strings.Contains(err.Error(), "complex") {
+		t.Errorf("%s error %q does not name the dtype", label, err)
+	}
+}
+
+// TestSecondOrderHelpersLeaveCallerGradients pins the internal reverse
+// passes of Hessian, HessianVectorProduct and MinimiseNewtonCG: the
+// objectives close over the caller's trainable tensors, and no path
+// (success or error) may mutate their accumulated gradients. AdjointODE
+// documents and implements the same guarantee.
+func TestSecondOrderHelpersLeaveCallerGradients(t *testing.T) {
+	thetaVals := []float64{2, 3}
+	presetVals := []float64{0.5, 0.25}
+	thetaArr, err := core.FromFloats(thetaVals, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	xArr, err := core.FromFloats([]float64{1, -1}, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	// Σ z_i²·θ_i: the Hessian is diag(2θ) = diag(4, 6).
+	weighted := func(theta *Tensor) func(*Tensor) (*Tensor, error) {
+		return func(z *Tensor) (*Tensor, error) {
+			sq, err := z.Mul(z)
+			if err != nil {
+				return nil, err
+			}
+			w, err := sq.Mul(theta)
+			if err != nil {
+				return nil, err
+			}
+			return w.Sum()
+		}
+	}
+	// Σ θ_i²: independent of the probes, the disconnected-objective case.
+	constant := func(theta *Tensor) func(*Tensor) (*Tensor, error) {
+		return func(*Tensor) (*Tensor, error) {
+			sq, err := theta.Mul(theta)
+			if err != nil {
+				return nil, err
+			}
+			return sq.Sum()
+		}
+	}
+
+	t.Run("hessian success", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		h, err := Hessian(weighted(theta), FromArray(xArr, false), HessianOptions{})
+		if err != nil {
+			t.Fatalf("Hessian: %v", err)
+		}
+		for i := range 2 {
+			if got, want := h.FloatAt(i*2+i), 2*thetaVals[i]; math.Abs(got-want) > 1e-8 {
+				t.Errorf("Hessian diagonal[%d] = %v, want %v", i, got, want)
+			}
+		}
+		requireGradUntouched(t, "Hessian", theta, preset, presetBits)
+	})
+
+	t.Run("hessian error", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		if _, err := Hessian(constant(theta), FromArray(xArr, false), HessianOptions{}); err == nil {
+			t.Fatal("expected the disconnected-objective error")
+		}
+		requireGradUntouched(t, "Hessian error path", theta, preset, presetBits)
+	})
+
+	t.Run("hvp success", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		v, err := core.FromFloats([]float64{1, 0.5}, 2)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		hv, err := HessianVectorProduct(weighted(theta), FromArray(xArr, false), FromArray(v, false), HessianOptions{})
+		if err != nil {
+			t.Fatalf("HessianVectorProduct: %v", err)
+		}
+		for i := range 2 {
+			if got, want := hv.FloatAt(i), 2*thetaVals[i]*v.FloatAt(i); math.Abs(got-want) > 1e-6 {
+				t.Errorf("H·v[%d] = %v, want %v", i, got, want)
+			}
+		}
+		requireGradUntouched(t, "HessianVectorProduct", theta, preset, presetBits)
+	})
+
+	t.Run("hvp error", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		v, err := core.FromFloats([]float64{1, 0.5}, 2)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		if _, err := HessianVectorProduct(constant(theta), FromArray(xArr, false), FromArray(v, false), HessianOptions{}); err == nil {
+			t.Fatal("expected the disconnected-objective error")
+		}
+		requireGradUntouched(t, "HessianVectorProduct error path", theta, preset, presetBits)
+	})
+
+	t.Run("newtoncg success", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		// Σ (z − θ)² minimises at z = θ, where the gradient vanishes.
+		objective := func(z *Tensor) (*Tensor, error) {
+			d, err := z.Sub(theta)
+			if err != nil {
+				return nil, err
+			}
+			sq, err := d.Mul(d)
+			if err != nil {
+				return nil, err
+			}
+			return sq.Sum()
+		}
+		x0, err := core.FromFloats([]float64{0, 0}, 2)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		got, fv, err := MinimiseNewtonCG(objective, x0, NewtonCGOptions{})
+		if err != nil {
+			t.Fatalf("MinimiseNewtonCG: %v", err)
+		}
+		for i := range 2 {
+			if math.Abs(got.FloatAt(i)-thetaVals[i]) > 1e-8 {
+				t.Errorf("minimiser[%d] = %v, want %v", i, got.FloatAt(i), thetaVals[i])
+			}
+		}
+		if math.Abs(fv) > 1e-12 {
+			t.Errorf("value = %v, want 0", fv)
+		}
+		requireGradUntouched(t, "MinimiseNewtonCG", theta, preset, presetBits)
+	})
+
+	t.Run("newtoncg error", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, presetBits := presetGrad(t, theta, presetVals)
+		x0, err := core.FromFloats([]float64{1}, 1)
+		if err != nil {
+			t.Fatalf("FromFloats: %v", err)
+		}
+		if _, _, err := MinimiseNewtonCG(constant(theta), x0, NewtonCGOptions{}); err == nil {
+			t.Fatal("expected the disconnected-objective error")
+		}
+		requireGradUntouched(t, "MinimiseNewtonCG error path", theta, preset, presetBits)
+	})
+}
+
+// presetGrad installs vals as theta's accumulated gradient and
+// returns the array the caller set together with a byte-exact snapshot
+// of its payload, so the check below can prove both that the very array
+// survived and that nothing wrote through it.
+func presetGrad(t *testing.T, theta *Tensor, vals []float64) (*core.Array, []uint64) {
+	t.Helper()
+	g, err := core.FromFloats(vals, theta.Data().Len())
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	theta.SetGrad(g)
+	return g, gradBits(g)
+}
+
+// gradBits copies the raw float bits of a gradient array.
+func gradBits(a *core.Array) []uint64 {
+	fs := a.RawFloats()
+	out := make([]uint64, len(fs))
+	for i, v := range fs {
+		out[i] = math.Float64bits(v)
+	}
+	return out
+}
+
+// requireGradUntouched requires the preset gradient to be the very
+// array that was set and to hold the exact bits the snapshot captured:
+// the helper under test must not write into it, replace it or clear it.
+func requireGradUntouched(t *testing.T, label string, theta *Tensor, want *core.Array, before []uint64) {
+	t.Helper()
+	got := theta.Grad()
+	if got == nil {
+		t.Errorf("%s: the caller's gradient was cleared", label)
+		return
+	}
+	if got != want {
+		t.Errorf("%s: the caller's gradient array was replaced", label)
+	}
+	after := gradBits(got)
+	if len(after) != len(before) {
+		t.Errorf("%s: the caller's gradient changed length, %d to %d", label, len(before), len(after))
+		return
+	}
+	for i := range before {
+		if after[i] != before[i] {
+			t.Errorf("%s: gradient[%d] = %v, want the preset bits of %v", label, i,
+				got.FloatAt(i), math.Float64frombits(before[i]))
+		}
+	}
+}
+
+// TestBroadcastToRefusesIntOperands pins the dtype gate on BroadcastTo,
+// the one shape op that used to accept an int tensor and record a graph
+// node for it. The shape is validated first, so an impossible target
+// keeps reporting the mismatch (the probe below asserts the same
+// order).
+func TestBroadcastToRefusesIntOperands(t *testing.T) {
+	src, err := core.FromInts([]int64{1, 2, 3}, 3)
+	if err != nil {
+		t.Fatalf("FromInts: %v", err)
+	}
+	if out, err := FromArray(src, true).BroadcastTo(2, 3); err == nil {
+		t.Errorf("BroadcastTo accepted an int tensor: shape %v dtype %s",
+			out.Data().Shape(), out.Data().Dtype())
+	} else if !strings.Contains(err.Error(), "needs a float") {
+		t.Errorf("int refusal = %v, want the dtype error", err)
+	}
+	one, err := core.FromInts([]int64{7}, 1)
+	if err != nil {
+		t.Fatalf("FromInts: %v", err)
+	}
+	if _, err := FromArray(one, true).BroadcastTo(2, 3); err == nil {
+		t.Error("BroadcastTo accepted an int (1,) tensor")
+	}
+	square, err := core.FromInts([]int64{1, 2, 3, 4}, 2, 2)
+	if err != nil {
+		t.Fatalf("FromInts: %v", err)
+	}
+	if _, err := FromArray(square, true).BroadcastTo(2, 3); err == nil {
+		t.Error("an impossible broadcast target was accepted")
+	} else if !strings.Contains(err.Error(), "cannot broadcast") {
+		t.Errorf("impossible int broadcast = %v, want the shape refusal", err)
+	}
+
+	// A float tensor still broadcasts and differentiates; a complex one
+	// remains inside the accepted dtypes.
+	xf, err := core.FromFloats([]float64{1, 2, 3}, 3)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	xt := FromArray(xf, true)
+	y, err := xt.BroadcastTo(2, 3)
+	if err != nil {
+		t.Fatalf("float BroadcastTo: %v", err)
+	}
+	loss, err := y.Sum()
+	if err != nil {
+		t.Fatalf("Sum: %v", err)
+	}
+	if err := loss.Backward(); err != nil {
+		t.Fatalf("Backward: %v", err)
+	}
+	for i := range xt.Grad().Len() {
+		if got := xt.Grad().FloatAt(i); got != 2 {
+			t.Errorf("float broadcast gradient[%d] = %v, want 2", i, got)
+		}
+	}
+
+	zc, err := core.FromComplexes([]complex128{1 + 1i, 2 - 0.5i}, 2)
+	if err != nil {
+		t.Fatalf("FromComplexes: %v", err)
+	}
+	zt := FromArray(zc, true)
+	cz, err := zt.BroadcastTo(3, 2)
+	if err != nil {
+		t.Fatalf("complex BroadcastTo: %v", err)
+	}
+	csq, err := cz.Abs2()
+	if err != nil {
+		t.Fatalf("Abs2: %v", err)
+	}
+	closs, err := csq.Sum()
+	if err != nil {
+		t.Fatalf("Sum: %v", err)
+	}
+	if err := closs.Backward(); err != nil {
+		t.Fatalf("Backward: %v", err)
+	}
+	// L = Σ|broadcast(z)|² = 3·Σ|z|², so dL/dz̄ = 3z.
+	for i := range 2 {
+		if got, want := zt.Grad().ComplexAt(i), 3*zc.ComplexAt(i); got != want {
+			t.Errorf("complex broadcast gradient[%d] = %v, want %v", i, got, want)
+		}
+	}
+}
+
+// TestComplexMeanOfEmptyIsAnError pins the degenerate reduction: the
+// complex branch of Mean divides by the element count, so an empty
+// tensor used to answer 0/0 = NaN while the real branch errors loudly.
+func TestComplexMeanOfEmptyIsAnError(t *testing.T) {
+	ec, err := core.FromComplexes([]complex128{}, 0)
+	if err != nil {
+		t.Fatalf("FromComplexes: %v", err)
+	}
+	m, cerr := FromArray(ec, true).Mean()
+	if cerr == nil {
+		t.Fatalf("complex Mean of an empty tensor returned %v instead of an error", m.Data().ComplexAt(0))
+	}
+	if !strings.Contains(cerr.Error(), "empty array has no mean") {
+		t.Errorf("complex Mean error = %v, want the empty-reduction refusal", cerr)
+	}
+	er, err := core.FromFloats([]float64{}, 0)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	_, rerr := FromArray(er, true).Mean()
+	if rerr == nil {
+		t.Fatal("real Mean of an empty tensor was accepted")
+	}
+	// The two dtypes answer with one message.
+	if rerr.Error() != cerr.Error() {
+		t.Errorf("messages disagree: real %q, complex %q", rerr, cerr)
+	}
+
+	// A non-empty complex mean still reduces and differentiates.
+	z, err := core.FromComplexes([]complex128{1 + 1i, 3 - 1i}, 2)
+	if err != nil {
+		t.Fatalf("FromComplexes: %v", err)
+	}
+	zt := FromArray(z, true)
+	mean, err := zt.Mean()
+	if err != nil {
+		t.Fatalf("Mean: %v", err)
+	}
+	if got, want := mean.Data().ComplexAt(0), complex(2, 0); got != want {
+		t.Errorf("complex mean = %v, want %v", got, want)
+	}
+	loss, err := mean.Abs2()
+	if err != nil {
+		t.Fatalf("Abs2: %v", err)
+	}
+	if err := loss.Backward(); err != nil {
+		t.Fatalf("Backward: %v", err)
+	}
+	// L = |mean|², so dL/dz̄ = mean/2 per element.
+	for i := range 2 {
+		if got, want := zt.Grad().ComplexAt(i), complex(1, 0); got != want {
+			t.Errorf("gradient[%d] = %v, want %v", i, got, want)
+		}
+	}
+}
+
+// TestAbs2KeepsFloat32Width pins Abs2's real branch: a float32 operand
+// squares in float64 and stays float32, exactly as Pow(2) does,
+// instead of promoting the forward result to float64.
+func TestAbs2KeepsFloat32Width(t *testing.T) {
+	vals := []float32{1.3, -2.7, 0.5}
+	a, err := core.FromFloat32s(vals, 3)
+	if err != nil {
+		t.Fatalf("FromFloat32s: %v", err)
+	}
+	xt := FromArray(a, true)
+	sq, err := xt.Abs2()
+	if err != nil {
+		t.Fatalf("Abs2: %v", err)
+	}
+	if got := sq.Data().Dtype(); got != core.Float32 {
+		t.Errorf("float32 Abs2 output dtype = %s, want float32", got)
+	}
+	pw, err := xt.Pow(2)
+	if err != nil {
+		t.Fatalf("Pow: %v", err)
+	}
+	if got := pw.Data().Dtype(); got != core.Float32 {
+		t.Errorf("float32 Pow(2) output dtype = %s, want float32", got)
+	}
+	for i, v := range vals {
+		want := float32(float64(v) * float64(v))
+		if got := sq.Data().FloatAt(i); got != float64(want) {
+			t.Errorf("Abs2[%d] = %v, want the once-rounded %v", i, got, want)
+		}
+		if got := pw.Data().FloatAt(i); got != float64(want) {
+			t.Errorf("Pow(2)[%d] = %v, want %v", i, got, want)
+		}
+	}
+
+	// L = Σx² has dL/dx = 2x, and the float32 leaf keeps its width.
+	loss, err := sq.Sum()
+	if err != nil {
+		t.Fatalf("Sum: %v", err)
+	}
+	if err := loss.Backward(); err != nil {
+		t.Fatalf("Backward: %v", err)
+	}
+	g := xt.Grad()
+	if got := g.Dtype(); got != core.Float32 {
+		t.Errorf("float32 leaf gradient dtype = %s, want float32", got)
+	}
+	for i, v := range vals {
+		if got, want := g.FloatAt(i), 2*float64(v); math.Abs(got-want) > 1e-6 {
+			t.Errorf("gradient[%d] = %v, want %v", i, got, want)
+		}
+	}
+
+	// A float64 operand is untouched: float64 out, exact squares.
+	af, err := core.FromFloats([]float64{1.3, -2.7}, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	sqf, err := FromArray(af, true).Abs2()
+	if err != nil {
+		t.Fatalf("Abs2: %v", err)
+	}
+	if got := sqf.Data().Dtype(); got != core.Float {
+		t.Errorf("float64 Abs2 output dtype = %s, want float", got)
+	}
+	for i := range 2 {
+		if got, want := sqf.Data().FloatAt(i), af.FloatAt(i)*af.FloatAt(i); got != want {
+			t.Errorf("float64 Abs2[%d] = %v, want %v", i, got, want)
+		}
+	}
+}
+
+// TestNewtonCGConstantObjectiveHitsDisconnectedGuard covers
+// MinimiseNewtonCG's g == nil guard, which the suite missed: its
+// "constant objective" case slices the point itself, so a gradient
+// exists and the guard never fires. A constant built from an
+// independent graduated tensor (the disconnected-objective pattern) has no path to
+// the starting point at all.
+func TestNewtonCGConstantObjectiveHitsDisconnectedGuard(t *testing.T) {
+	c, err := FromFloat64s([]float64{3}, true, 1)
+	if err != nil {
+		t.Fatalf("FromFloat64s: %v", err)
+	}
+	constant := func(*Tensor) (*Tensor, error) { return c.Mul(c) }
+	x0, err := core.FromFloats([]float64{1}, 1)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	if _, _, err := MinimiseNewtonCG(constant, x0, NewtonCGOptions{MaxIterations: 2}); err == nil {
+		t.Fatal("a genuinely constant objective minimised without error")
+	} else if !strings.Contains(err.Error(), "does not depend on the starting point") {
+		t.Fatalf("error = %v, want the disconnected-graph refusal", err)
+	}
+}
+
+// TestSampleHMCLeavesCallerGradients pins SampleHMC's internal gradient
+// evaluations: one leapfrog step runs one reverse pass, so a density
+// closing over a graduated tensor used to add a contribution per step
+// and one per proposal, on the success and the error path alike. The
+// pass commits nothing, so the closed-over tensor keeps its accumulated
+// gradient bit for bit, the guarantee Hessian, HessianVectorProduct,
+// MinimiseNewtonCG and AdjointODE document.
+func TestSampleHMCLeavesCallerGradients(t *testing.T) {
+	thetaVals := []float64{2, 1.5}
+	presetVals := []float64{0.5, -0.25}
+	thetaArr, err := core.FromFloats(thetaVals, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	q0, err := core.FromFloats([]float64{0.5, -0.5}, 2)
+	if err != nil {
+		t.Fatalf("FromFloats: %v", err)
+	}
+	// log π(q) = −½·Σ θ_i q_i²: a density that closes over a graduated
+	// tensor the caller owns, so any committed pass shows up in θ.
+	density := func(theta *Tensor) func(*Tensor) (*Tensor, error) {
+		return func(q *Tensor) (*Tensor, error) {
+			sq, err := q.Mul(q)
+			if err != nil {
+				return nil, err
+			}
+			w, err := sq.Mul(theta)
+			if err != nil {
+				return nil, err
+			}
+			s, err := w.Sum()
+			if err != nil {
+				return nil, err
+			}
+			return s.Scale(-0.5)
+		}
+	}
+
+	t.Run("success", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, bits := presetGrad(t, theta, presetVals)
+		calls := 0
+		counted := func(q *Tensor) (*Tensor, error) {
+			calls++
+			return density(theta)(q)
+		}
+		samples, err := SampleHMC(counted, q0,
+			HMCOptions{Step: 0.1, Steps: 2, Samples: 2, Thin: 1, Seed: 11})
+		if err != nil {
+			t.Fatalf("SampleHMC: %v", err)
+		}
+		if got := samples.Shape(); got[0] != 2 || got[1] != 2 {
+			t.Fatalf("samples shape %v, want [2 2]", got)
+		}
+		// One evaluation at q0 plus one per leapfrog step: the run really
+		// did differentiate the density several times.
+		if calls < 3 {
+			t.Fatalf("SampleHMC ran %d gradient evaluations, want at least 3", calls)
+		}
+		requireGradUntouched(t, "SampleHMC", theta, preset, bits)
+	})
+
+	t.Run("midChainError", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, bits := presetGrad(t, theta, presetVals)
+		// The first evaluation at q0 succeeds, so a reverse pass has run;
+		// every proposal then reports the state as outside the support,
+		// which rejects the trajectory instead of aborting the run.
+		calls, refused := 0, 0
+		failing := func(q *Tensor) (*Tensor, error) {
+			calls++
+			if calls > 1 {
+				refused++
+				return nil, errf("outside the support")
+			}
+			return density(theta)(q)
+		}
+		samples, err := SampleHMC(failing, q0,
+			HMCOptions{Step: 0.1, Steps: 3, Samples: 1, Thin: 1, Seed: 12})
+		if err != nil {
+			t.Fatalf("SampleHMC: %v", err)
+		}
+		if got := samples.Shape(); got[0] != 1 || got[1] != 2 {
+			t.Fatalf("samples shape %v, want [1 2]", got)
+		}
+		if refused == 0 {
+			t.Fatal("the density was never forced to fail mid-chain")
+		}
+		requireGradUntouched(t, "SampleHMC mid-chain error", theta, preset, bits)
+	})
+
+	t.Run("startError", func(t *testing.T) {
+		theta := FromArray(thetaArr, true)
+		preset, bits := presetGrad(t, theta, presetVals)
+		fails := func(*Tensor) (*Tensor, error) { return nil, errf("no density at the start") }
+		if _, err := SampleHMC(fails, q0,
+			HMCOptions{Step: 0.1, Steps: 2, Samples: 1, Thin: 1, Seed: 13}); err == nil {
+			t.Fatal("expected the start-time density error to be fatal")
+		}
+		requireGradUntouched(t, "SampleHMC start error", theta, preset, bits)
+	})
+}