第4章：计算图——GoMLX 的核心数据结构

---

4.1 从一段代码到一张图

先看一段你已经能写的 GoMLX 代码：

func myPoly(x *graph.Node) *graph.Node {
    x2 := graph.Square(x)
    x4 := graph.Square(x2)
    return graph.MulScalar(x4, 3.0)
}

这段代码翻译成数学公式是 f(x) = 3x⁴。但在 GoMLX 内部，它首先被翻译成一张这样的图：

   [x]
    |
  [Square]  ← 节点 1: x²
    |
   [x2]
    |
  [Square]  ← 节点 2: (x²)² = x⁴
    |
   [x4] ─────┐
             │
  [Const 3.0]┤
             │
          [MulScalar]  ← 节点 3: x⁴ * 3
             |
          [output]

这就是计算图（Computation Graph）。它是一张有向无环图（DAG，Directed Acyclic Graph）：

节点（Node）：每一个运算的结果—— Square、MulScalar、输入占位符 x
边（Edge）：数据依赖—— x2 依赖 x，x4 依赖 x2，输出依赖 x4 和常量 3.0
无环：数据只向前流动，没有循环。循环是训练循环的事，不在这张图里。

这张图不包含任何数值。x 的值、平方的结果，这些都不在图里。图只记录关系和形状。它像一个模板——编译器拿着这个模板，再结合你实际输入的 x 的值，生成可执行的机器码。

---

4.2 `*graph.Graph`：整张图的容器

在 GoMLX 中，*graph.Graph 是整张计算图的持有者：

// 函数式构图时，GoMLX 自动创建和管理 Graph
exec := graph.MustNewExec(backend, myFn)
// Graph 由 Exec 内部持有，你通常不直接操作它

// 手动构图时，你显式创建 Graph
g := graph.NewGraph(backend, "my-graph-name")

一张 *Graph 持有以下信息：

属性	说明
名称	调试用标识（如 `"my-add-graph"`）
后端引用	指向 `backends.Backend`，在编译阶段用到
节点列表	所有被创建的 Node，按构建顺序记录
参数节点	标记为"输入"的节点——即图的入口
输出节点	编译时指定的输出节点——即图的出口

在图构建过程中，每当你调用 graph.Add、graph.Square 等 API，底层的 *Graph 就把新节点追加到节点列表里。这些节点之间的引用关系形成了 DAG 的边。

大部分时候你不需要感知 *Graph 的存在。函数式构图时 GoMLX 在背后自动管理一切。但当你需要条件分支（比如根据输入构建不同分支的图）或共享子图时，理解 *Graph 就变得重要了。

---

4.3 `*graph.Node`：每一个运算的元数据

一个 *graph.Node 不只是"一个数字的占位符"。它是一个结构体，携带以下信息：

Node {
    graph:     *Graph         // 所属的图
    opType:    OpType         // 运算类型：Add / Square / MulScalar ...
    inputs:    []*Node        // 输入节点（边的来源）
    shape:     shapes.Shape   // 输出形状和数据类型
    id:        NodeId         // 在图中唯一的数字 ID
}

用你熟悉的 Go 思维来理解：graph.Add(a, b) 等价于创建了一个这样的结构体：

// 伪代码——真实实现细节不同
func Add(a, b *Node) *Node {
    n := &Node{
        graph:  a.graph,     // 和输入节点属于同一张图
        opType: AddOp,
        inputs: []*Node{a, b},
        shape:  broadcastShapes(a.shape, b.shape), // 形状自动推导
    }
    a.graph.addNode(n)       // 把新节点注册到图上
    return n
}

这里有几个关键推论。

Node 不存值，只存关系。 在 MustExec 之前，任何 Node 的"值"都是未定义的。图是一张蓝图，不是一张数据表。

形状在构建时就确定了。 broadcastShapes 在图构建阶段运行，所以形状错误在编译前就能发现——这也是 GoMLX 类型安全的一个重要维度。

同一个 Node 可以被多次引用。 如果你写了 graph.Add(sameNode, anotherNode)，图中的 sameNode 只算一次。GoMLX 在编译时通过"公共子表达式消除"（CSE）确保这一点。你可以大胆复用同一个 *Node，不必担心重复计算。

---

4.4 图构建的底层：两种 API 背后的同一张图

第 2 章简单提过两种构图方式。现在从底层看它们各自在做什么。

函数式构图：自动追踪

exec := graph.MustNewExec(backend, func(x *graph.Node) *graph.Node {
    return graph.Square(x)
})

MustNewExec 内部做的事： 1. 调用 backend.NewGraph() 创建一个新 *Graph 2. 在图中创建一个参数节点（代表 x） 3. 调用你的函数，传入参数节点；GoMLX 的 graph.Square 等 API 在每次调用时把新节点添加到图上 4. 你的函数返回的节点（或节点列表）被标记为图的输出 5. 后端拿到完整的 *Graph，开始编译

你用 graph.Square(x) 时，GoMLX 在内部读取当前活跃的图上下文，找到对应的 *Graph，往上追加节点。这一切都是自动的——你写代码的感觉跟写普通 Go 函数几乎没有区别。

手动构图：精确控制

g := graph.NewGraph(backend, "my-graph")

x := graph.Parameter(g, "x", shapes.Make(dtypes.Float64))
y := graph.Parameter(g, "y", shapes.Make(dtypes.Float64))

z := graph.Add(x, y)

err := g.Compile(z)  // 指定 z 为输出，编译图
if err != nil {
    panic(err)
}

results := g.Run(3.0, 4.0)
fmt.Println(results[0].Value()) // 7

手动构图的每步操作都是显式的。你创建 *Graph，创建每个参数节点并命名，逐个添加运算节点，手动指定输出节点并调用 g.Compile()。编译成功后，用 g.Run() 直接传入输入值执行。

函数式构图做不到的几件事，手动构图可以： 1. 命名节点：你可以给每个参数节点起个有意义的英文名，调试时堆栈更有可读性 2. 条件分支：你可以在 if 分支里构建不同的子图结构 3. 共享子图：你可以手动管理子图复用，避免重复构图

怎么选？

场景	推荐方式
日常模型定义	函数式构图
需要动态分支（不同超参构建不同结构）	手动构图
需要精细命名（调试/日志）	手动构图
子图复用（同一个特征提取器在多个地方被调用）	手动构图

---

4.5 在图上行走：形状推导、遍历和依赖分析

有了整张 DAG，GoMLX 可以做几件有趣的事情。

形状推导（Shape Inference）

你写 graph.Add(a, b) 时，不需要手动指定结果的形状——GoMLX 从 a.Shape() 和 b.Shape() 自动推导。推导规则包括：

逐元素运算（Add、Mul、Square 等）：输出形状 = 输入形状的广播结果
矩阵乘法（DotGeneral）：输出形状 = [M, K] × [K, N] = [M, N]
缩减运算（ReduceSum、ReduceMean）：输出形状 = 去掉缩减轴后的输入形状

这些推导在图构建阶段就完成了。形状不匹配的错误不会等到运行时——在 MustNewExec 返回之前就会 panic。

拓扑遍历（Topological Traversal）

图编译时，GoMLX 从输出节点出发，按拓扑顺序（依赖在前）遍历所有节点。这是任何 DAG 引擎的标准算法。

公共子表达式消除（CSE）

如果你写了相同输入和运算的节点，GoMLX 的编译器会把它们合并成一个节点——只算一次。

---

4.6 图调试：看到你看不见的

方式一：节点日志（SetLogged）

func debugGraph(x *graph.Node) *graph.Node {
    x2 := graph.Square(x)
    x2.SetLogged("x²")
    x4 := graph.Square(x2)
    x4.SetLogged("x⁴")
    return graph.MulScalar(x4, 3.0)
}

SetLogged 标记的节点，在每次 MustExec 后自动打印其值。

方式二：图结构输出

fmt.Println(g.String())
// 输出类似：
// Graph "debug": 4 nodes, 1 parameters
//   #0	Parameter "x" → Float64[]
//   #1	Square(#0) → Float64[]
//   #2	Square(#1) → Float64[]
//   #3	MulScalar(#2, 3) → Float64[]

---

4.7 计算图与训练：变量怎么融进去

在实际模型中，你有一堆可学参数（权重和偏置）。GoMLX 通过变量节点（Variable Node）通过 context.Context 创建，在图中表现为一个"可读可写的 *Node"：

ctx := context.New()
ctx = ctx.In("my_model")
wVar := ctx.VariableWithValue("w", float64(0.5))

g := graph.NewGraph(backend, "var-demo")
x := graph.Parameter(g, "x", shapes.Make(dtypes.Float64))

w := wVar.ValueGraph(g)  // 当前值为 0.5
y := graph.Mul(x, w)     // y = x * 0.5

g.Compile(y)
results := g.Run(3.0)
fmt.Println(results[0].Value()) // 3.0 * 0.5 = 1.5

在训练循环里，wVar.SetValueGraph(newW) 更新变量在图中的值。context.Context 和 context.Variable 构成了 GoMLX 的参数管理体系。

---

4.8 收束：从数据结构到系统设计

一张图总结本章的核心认知：

GoMLX 的计算图 = DAG（有向无环图）
    │
    ├── 节点 = 运算（Add, Square, MulScalar ...）+ 形状信息
    ├── 边 = 数据依赖（topological order 保证编译正确）
    ├── 构建方式 = 函数式（自动追踪）或手动（精确控制）
    ├── 编译期优化 = 形状推导 + CSE + 算子融合（第 5 章）
    └── 动态状态 = 变量节点（context.Variable）让纯函数图拥有记忆

---

> 作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》 > 上一章：第3章 GoMLX的世界观 > 下一章预告：第5章 XLA与JIT编译——从图到机器码的全过程