feat: initial release

Assisted-by: GLM 5.3 Flash

grad/tensor_test.go

@@ -0,0 +1,343 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package grad
+
+import (
+	"math"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+	"strings"
+	"testing"
+)
+
+func mustTensor(t *testing.T, vals []float64, shape ...int) *Tensor {
+	t.Helper()
+	tt, err := FromFloat64s(vals, true, shape...)
+	if err != nil {
+		t.Fatalf("FromFloat64s(%v, %v): %v", vals, shape, err)
+	}
+	return tt
+}
+
+func TestAutogradBasicChain(t *testing.T) {
+	x := mustTensor(t, []float64{2, 3}, 2)
+	y := mustTensor(t, []float64{4, 5}, 2)
+	z, err := x.Add(y)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sq, err := z.Mul(z)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s, err := sq.Sum()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	// d/dx (x+y)^2 summed = 2(x+y); at x=2: 12, at x=3: 16.
+	if gx := x.Grad().FloatAt(0); math.Abs(gx-12) > 1e-9 {
+		t.Errorf("grad x[0]: got %v, want 12", gx)
+	}
+	if gx := x.Grad().FloatAt(1); math.Abs(gx-16) > 1e-9 {
+		t.Errorf("grad x[1]: got %v, want 16", gx)
+	}
+	// y's gradient matches x's, symmetric in the sum.
+	if gy := y.Grad().FloatAt(0); math.Abs(gy-12) > 1e-9 {
+		t.Errorf("grad y[0]: got %v, want 12", gy)
+	}
+}
+
+func TestAutogradMatMul(t *testing.T) {
+	a := mustTensor(t, []float64{1, 2, 3, 4}, 2, 2)
+	b := mustTensor(t, []float64{5, 6, 7, 8}, 2, 2)
+	p, err := a.MatMul(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s, err := p.Sum()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	// d/dA sum(A·B) = J·Bᵀ; Bᵀ = [[5,7],[6,8]], so J·Bᵀ =
+	// [[11,15],[11,15]] (each row is the column sums of Bᵀ).
+	wantA := []float64{11, 15, 11, 15}
+	for i := range 4 {
+		if g := a.Grad().FloatAt(i); math.Abs(g-wantA[i]) > 1e-9 {
+			t.Errorf("grad A[%d]: got %v, want %v", i, g, wantA[i])
+		}
+	}
+	// d/dB sum(A·B) = Aᵀ·J; Aᵀ = [[1,3],[2,4]], row sums: 4, 6, so
+	// Aᵀ·J = [[4,4],[6,6]].
+	wantB := []float64{4, 4, 6, 6}
+	for i := range 4 {
+		if g := b.Grad().FloatAt(i); math.Abs(g-wantB[i]) > 1e-9 {
+			t.Errorf("grad B[%d]: got %v, want %v", i, g, wantB[i])
+		}
+	}
+}
+
+func TestAutogradMatMulVector(t *testing.T) {
+	// 2-D × 1-D: y = A·x.
+	a := mustTensor(t, []float64{1, 2, 3, 4}, 2, 2)
+	x := mustTensor(t, []float64{2, 3}, 2)
+	y, err := a.MatMul(x)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s, err := y.Sum()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	// grad x = Aᵀ·1 = column sums of A: 4, 6.
+	if g := x.Grad().FloatAt(0); math.Abs(g-4) > 1e-9 {
+		t.Errorf("grad x[0]: got %v, want 4", g)
+	}
+	if g := x.Grad().FloatAt(1); math.Abs(g-6) > 1e-9 {
+		t.Errorf("grad x[1]: got %v, want 6", g)
+	}
+	// grad A = outer(1, x): [[2,3],[2,3]].
+	if g := a.Grad().FloatAt(2); math.Abs(g-2) > 1e-9 {
+		t.Errorf("grad A[2]: got %v, want 2", g)
+	}
+}
+
+func TestAutogradActivations(t *testing.T) {
+	// Sigmoid at 0: σ(0)=0.5, σ' = 0.25.
+	x3 := mustTensor(t, []float64{0}, 1)
+	sg, _ := x3.Sigmoid()
+	if err := sg.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := x3.Grad().FloatAt(0); math.Abs(g-0.25) > 1e-9 {
+		t.Errorf("Sigmoid grad at 0: got %v, want 0.25", g)
+	}
+
+	// Exp and Log compose to identity: grad log(exp(x)) = 1.
+	x4 := mustTensor(t, []float64{2}, 1)
+	e, _ := x4.Exp()
+	l, _ := e.Log()
+	if err := l.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := x4.Grad().FloatAt(0); math.Abs(g-1) > 1e-9 {
+		t.Errorf("log(exp) grad: got %v, want 1", g)
+	}
+}
+
+func TestAutogradGradientAccumulation(t *testing.T) {
+	x := mustTensor(t, []float64{1}, 1)
+	a, _ := x.Mul(x)
+	b, _ := x.Mul(x)
+	s, err := a.Add(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	// d/dx (x² + x²) at x=1 = 4.
+	if g := x.Grad().FloatAt(0); math.Abs(g-4) > 1e-9 {
+		t.Errorf("accumulated grad: got %v, want 4", g)
+	}
+	// A second Backward without ZeroGrad accumulates into the leaf.
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := x.Grad().FloatAt(0); math.Abs(g-8) > 1e-9 {
+		t.Errorf("accumulated grad after 2nd pass: got %v, want 8", g)
+	}
+	x.ZeroGrad()
+	if x.Grad() != nil {
+		t.Error("ZeroGrad did not clear the gradient")
+	}
+}
+
+func TestAutogradRejectsNonFloat(t *testing.T) {
+	i, err := core.FromInts([]int64{1, 2}, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	it := FromArray(i, true)
+	if _, err := it.Sum(); err == nil || !strings.Contains(err.Error(), "float") {
+		t.Errorf("int Sum: %v", err)
+	}
+	c, _ := core.FromComplexes([]complex128{1 + 2i}, 1)
+	ct := FromArray(c, true)
+	// Complex Exp is differentiable (the Wirtinger graph); the
+	// real-only kernels are the ones that must still refuse it.
+	if _, err := ct.Exp(); err != nil {
+		t.Errorf("complex Exp must differentiate: %v", err)
+	}
+	if _, err := ct.Log(); err == nil {
+		t.Error("complex Log must error")
+	}
+	if _, err := ct.Tanh(); err == nil {
+		t.Error("complex Tanh must error")
+	}
+}
+
+func TestAutogradDivTanhNeg(t *testing.T) {
+	// d/dx (x/y) at x=4, y=2 = 1/2.
+	x := mustTensor(t, []float64{4}, 1)
+	y := mustTensor(t, []float64{2}, 1)
+	q, err := x.Div(y)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := q.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := x.Grad().FloatAt(0); math.Abs(g-0.5) > 1e-9 {
+		t.Errorf("Div grad x: got %v, want 0.5", g)
+	}
+	// d/dy (x/y) at y=2 = -x/y² = -1.
+	if g := y.Grad().FloatAt(0); math.Abs(g+1) > 1e-9 {
+		t.Errorf("Div grad y: got %v, want -1", g)
+	}
+
+	// tanh'(0) = 1.
+	t0 := mustTensor(t, []float64{0}, 1)
+	th, _ := t0.Tanh()
+	if err := th.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := t0.Grad().FloatAt(0); math.Abs(g-1) > 1e-9 {
+		t.Errorf("Tanh grad at 0: got %v, want 1", g)
+	}
+
+	// d/dx (-x) = -1.
+	n := mustTensor(t, []float64{3}, 1)
+	neg, _ := n.Neg()
+	if err := neg.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	if g := n.Grad().FloatAt(0); g != -1 {
+		t.Errorf("Neg grad: got %v, want -1", g)
+	}
+
+	// Accessors and Mean grad.
+	m := mustTensor(t, []float64{1, 2, 3, 4}, 2, 2)
+	if m.Data() != m.Data() || m.RequiresGrad() != true {
+		t.Error("accessors wrong")
+	}
+	mean, err := m.Mean()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := mean.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	// d/dx mean(x) = 1/n = 1/4.
+	if g := m.Grad().FloatAt(0); math.Abs(g-0.25) > 1e-9 {
+		t.Errorf("Mean grad: got %v, want 0.25", g)
+	}
+	if m.Grad() == nil {
+		t.Error("Grad() must be non-nil after Backward")
+	}
+}
+
+func TestAutogradReshape(t *testing.T) {
+	x, err := FromFloat64s([]float64{1, 2, 3, 4}, true, 2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r, err := x.Reshape(4)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sumT, err := r.Sum()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := sumT.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	g, err := x.Grad().Elements[float64]()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range g {
+		if g[i] != 1 {
+			t.Fatalf("Reshape grad[%d]: %v, want 1", i, g[i])
+		}
+	}
+}
+
+func TestAutogradTransposeBackward(t *testing.T) {
+	x, err := FromFloat64s([]float64{1, 2, 3, 4}, true, 2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tr, err := x.Transpose()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s, err := tr.Sum()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Backward(); err != nil {
+		t.Fatal(err)
+	}
+	g, err := x.Grad().Elements[float64]()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := range g {
+		if g[i] != 1 {
+			t.Errorf("Transpose grad[%d]: %v, want 1", i, g[i])
+		}
+	}
+}
+
+// TestAutogradPowZeroGradient pins the exponent-0 backward: d/dx x⁰
+// is the zero gradient everywhere, including at x = 0 where the
+// chain rule would evaluate 0·∞ and produce NaN.
+func TestAutogradPowZeroGradient(t *testing.T) {
+	x := mustTensor(t, []float64{0, 2}, 2)
+	y, err := x.Pow(0)
+	if err != nil {
+		t.Fatalf("Pow(0): %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 2 {
+		g := x.Grad().FloatAt(i)
+		if math.IsNaN(g) || g != 0 {
+			t.Errorf("d/dx x⁰ at %g = %v, want exactly 0", x.Data().FloatAt(i), g)
+		}
+	}
+}
+
+// TestAutogradLeafBackwardAccumulates pins that Backward on a leaf
+// accumulates into the existing gradient like any other backward
+// pass, instead of overwriting it.
+func TestAutogradLeafBackwardAccumulates(t *testing.T) {
+	x := mustTensor(t, []float64{3}, 1)
+	if err := x.Backward(); err != nil {
+		t.Fatalf("Backward: %v", err)
+	}
+	if got := x.Grad().FloatAt(0); got != 1 {
+		t.Fatalf("first leaf Backward: grad %v, want 1", got)
+	}
+	if err := x.Backward(); err != nil {
+		t.Fatalf("second Backward: %v", err)
+	}
+	if got := x.Grad().FloatAt(0); got != 2 {
+		t.Fatalf("second leaf Backward: grad %v, want 2 (accumulated)", got)
+	}
+}