第6章：张量——数据的形状与流动

6.1 张量，不止是"多维数组"

大多数资料对张量的解释是"多维数组"。这个定义没错，但对你理解 GoMLX 中的张量几乎没有帮助。

在 GoMLX 里，张量是一个有形状的数据容器，同时存在于两个空间：

• 逻辑空间：你看它是什么——比如一个 [128, 512] 的 float32 矩阵
• 物理空间：它存在哪——是 Go 程序的内存（本地），还是 GPU 的显存（设备）

这两个空间的关系，决定了你的 GoMLX 代码有多快、会不会内存泄漏。

6.2 数据结构：Shape + DType = 张量的身份证

GoMLX 中，张量的"类型信息"由两个东西唯一确定：形状（shapes.Shape）和数据类型（dtypes.DType）。

import (
    "github.com/gomlx/gomlx/pkg/core/shapes"
    "github.com/gomlx/gomlx/pkg/core/dtypes"
)

shape1 := shapes.Make(dtypes.Float32, 128, 512)
// 含义：数据类型 = float32，维度 = [128, 512]，总元素数 = 65536

shapes.Shape 包含几个核心属性：

DType 的选择直接影响内存和精度：

6.3 本地张量 vs 设备张量：理解这层关系就理解了性能

GoMLX 的张量，核心设计是双面存储：同一个 *tensors.Tensor 对象，可以同时在 Go 内存（本地）和设备内存（GPU 显存）持有数据。

import "github.com/gomlx/gomlx/pkg/core/tensors"

t1 := tensors.FromValue([][]float32{
    {1.0, 2.0, 3.0},
    {4.0, 5.0, 6.0},
})
// t1 的本地数据可直接读写

// 当你把 t1 传给 exec.MustExec 时，GoMLX 自动做：
//   1. 如果设备上没有 t1 的副本 → 复制到设备（慢）
//   2. 后续调用 → 设备数据复用（快）
//   3. 如果你修改了 t1 的本地数据 → 设备副本失效，下次再复制
result := exec.MustExec1(t1)

这个过程的关键：

1. 自动同步：GoMLX 在你把张量传入 MustExec 时，自动检查设备上是否有最新的副本。
2. 同步不是免费的：从 Go 内存到 GPU 显存的复制，带宽约 5-12 GB/s（PCIe 3.0/4.0）。
3. 最佳实践：训练数据在 Dataset 层预加载到设备上，张量留在设备侧，只在需要可视化或保存 Checkpoint 时才 MaterializeLocal() 搬回 Go 一侧。

6.4 Finalize：GoMLX 的垃圾回收

Go 程序员信赖 GC。但在 GoMLX 里，Go 的 GC 不管设备内存。GPU 显存不在 Go 的堆上——GC 看不到它，也不会回收它。

这就是 FinalizeAll() 存在的理由：

for i := 0; i < 1000000; i++ {
    result := exec.MustExec1(x)
    // 此时 result 占用了设备内存（可能几百 KB）
    result.FinalizeAll() // ← 不调这行 = 内存泄漏！
}

FinalizeAll() 做了什么：

• 释放张量在设备上的内存（如果这是唯一引用）
• 释放张量在本地 Go 内存中的数据（如果这是唯一引用）
• 把张量标记为"已终结"——之后再访问它会 panic

什么时候调？

忘了调的后果不是程序立刻崩，而是 GPU 显存持续增长——直到某一天 OOM。

6.5 形状操作：reshape、broadcast、切片

Broadcast（广播）

func broadcastExample(g *graph.Graph) *graph.Node {
    scalar := graph.Scalar(g, dtypes.Float32, 2.0)
    vec := graph.Iota(g, shapes.Make(dtypes.Float32, 3, 1), 0)  // iotaAxis=0
    return graph.Add(scalar, vec)  // 自动广播
}

广播规则：每个维度，要么相等，要么其中一个是 1。这和 NumPy 的广播完全一致。

Reshape（变形）

original := graph.Iota(g, shapes.Make(dtypes.Float32, 2, 3, 4), 0) // [2,3,4]
flattened := graph.Reshape(original, 2, 12)                          // [2,12]

Reshape 不移动数据——只改变"怎么看"张量的维度。但总元素数必须不变。

Slice（切片）

graph.Slice 接受一个节点和一组由 graph.AxisRange 构建的规格：

full := graph.Parameter(g, "x", shapes.Make(dtypes.Float32, 128, 512))
first10 := graph.Slice(full, graph.AxisRange(0, 10), graph.AxisRange()) // [10, 512]

graph.AxisRange 的几种用法：

• AxisRange()：取该轴全部（等同于 :）
• AxisRange(i)：取单个索引 i
• AxisRange(start, end)：取 [start, end) 范围
• AxisRange(start, end, stride)：取 [start, end) 范围，步长为 stride

6.6 DType 转换：精度与性能的权衡

func dtypeExample(x *graph.Node) *graph.Node {
    float32x := graph.ConvertDType(x, dtypes.Float32) // Float64 → Float32
    result := graph.Square(float32x)
    return graph.ConvertDType(result, dtypes.Float64) // 回到 Float64
}

转换的典型场景：

6.7 张量的完整生命周期

                    张量生命周期
  
  1. 创建
     tensors.FromValue(data)  │  graph.Parameter(g, ...)
                              │  graph.Scalar(g, ...)
     ┌──────────┐
     │ 本地数据  │
     └──────────┘
          │ 传入 MustExec
          ▼
  2. 搬运到设备（隐式，自动）
     ┌──────────┐
     │ 设备数据  │  ← GPU 显存 / CPU 设备内存
     └──────────┘
          │ MustExec 返回结果
          ▼
  3. 结果张量
     ┌──────────┐
     │ 设备数据  │  ← 执行结果在设备上
     └──────────┘
          │
          ├──→ result.FinalizeAll() ← 必须调！
          │
          └──→ result.MaterializeLocal() ← 可选，搬回 Go
               ┌──────────┐
               │ 本地数据  │  ← 现在可以读 Value()
               └──────────┘

   ⚠ 忘了 FinalizeAll → 设备内存泄漏 → OOM

6.8 收束

本章核心就五条规则：

1. 张量 = 形状 + DType + 双面存储（本地 / 设备）
2. 本地 ↔ 设备的搬运是隐式的，但是昂贵的——批量操作、避免来回搬
3. FinalizeAll() 不是建议，是必须——Go 的 GC 不回收 GPU 显存
4. 形状操作（broadcast / reshape / slice）不搬数据，只改视角
5. Float32 是默认选择，Float16 用于推理加速，Float64 用于需要极限精度的场景

第 7 章——自动微分。那是全书最核心的技术章节：梯度到底是怎么精确算出来的。

作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》
上一章：第5章 XLA与JIT编译
下一章预告：第7章 自动微分——梯度是怎样"算"出来的