《GoMLX从入门到精通》附录A：GoMLX API 速查表

计算图核心

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

// 图定义与执行
graph.MustNewExec(backend, graphFn)   // 编译图
exec.MustExec(args...)                // 执行 → 多输出
exec.MustExec1(args...)               // 执行 → 单输出
graph.Gradient(output, nodes...)      // 自动微分

// 基础运算
graph.Add(a, b)                       // 加法
graph.Sub(a, b)                       // 减法
graph.Mul(a, b)                       // 乘法
graph.Square(x)                       // 平方
graph.MulScalar(x, scalar)            // 标量乘法
graph.AddScalar(x, scalar)            // 标量加法
graph.Scalar(g, dtype, value)         // 常数节点

// 形状操作
graph.Reshape(x, dims...)             // 变形
graph.Iota(g, shape, iotaAxis)        // 递增序列（3个参数）
graph.Slice(x, specs...)              // 切片，配合 AxisRange
graph.AxisRange()                     // 全轴 (:)
graph.AxisRange(start, end)           // [start, end)
graph.AxisRange(i)                    // 单索引
graph.ReduceAllSum(x)                 // 全缩减求和
graph.ReduceAllMean(x)                // 全缩减求均值
graph.ConvertDType(x, dtype)          // 类型转换

// 图信息
n.Graph()                             // 获取所属图
n.Shape()                             // 获取形状
n.SetLogged("msg")                    // 执行后打印值

上下文（context）

import "github.com/gomlx/gomlx/pkg/ml/context"

ctx := context.New()                  // 创建上下文（无参数）
ctx.SetParam(key, value)              // 设置超参数
ctx.In("scope")                       // 进入子作用域
ctx.GetVariableByScopeAndName(scope, name) // 查询变量

// 带 Context 的 Exec
context.MustNewExec(backend, ctx, ctxGraphFn)
exec.MustExec1(input)

Checkpoint

import "github.com/gomlx/gomlx/pkg/ml/context/checkpoints"

// 从目录加载/保存
handler, err := checkpoints.Build(ctx).Dir("path/to/ckpt").Keep(3).Done()
handler.Save()                        // 保存 Checkpoint

// 从 go:embed 加载
_, err = checkpoints.Build(ctx).FromEmbed(jsonStr, binBytes).Done()

// 训练循环钩子
train.EveryNSteps(loop, 100, "checkpointing", priority, handler.OnStepFn)

层（layers）

import "github.com/gomlx/gomlx/pkg/ml/layers"

layers.DenseWithBias(ctx, input, units)  // 全连接+偏置
layers.Dense(ctx, input, useBias, units) // 全连接（可关偏置）
layers.DropoutStatic(ctx, input, 0.5)    // Dropout（静态 rate）
layers.Dropout(ctx, input, rateNode)     // Dropout（rate 为 *Node）

// 批归一化——在子包 batchnorm 中
import "github.com/gomlx/gomlx/pkg/ml/layers/batchnorm"
// 参阅子包文档

// 层归一化——返回 Builder，需调 .Done()
layers.LayerNormalization(ctx, input, axis...).Done()

// 卷积
layers.Convolution(ctx, input).
    Filters(n).KernelSize(h, w).Strides(s).PadSame().Done()

// 最大池化——在 graph 包中
graph.MaxPool(input).Window(k).Strides(s).Done()

// 注意力——返回 Builder，需调 .Done()
attention.MultiHeadAttention(ctx, q, k, v, numHeads, headDim).Done()

// 激活
import "github.com/gomlx/gomlx/pkg/ml/layers/activations"
activations.Relu(x)
activations.Sigmoid(x)
activations.Tanh(x)
activations.Gelu(x)
activations.Softmax(x)

训练（train）

import "github.com/gomlx/gomlx/pkg/ml/train"

trainer := train.NewTrainer(backend, ctx, modelFn, lossFn, optimizer, metrics, regularizers)
loop := train.NewLoop(trainer)
loop.RunSteps(ds, numSteps)
loop.Trainer.Eval(ds)

// 钩子——注意 priority 参数
train.EveryNSteps(loop, n, name, priority, fn) // 每 N 步触发

// 优化器
import "github.com/gomlx/gomlx/pkg/ml/train/optimizers"
optimizers.StochasticGradientDescent().WithLearningRate(0.01).Done()
optimizers.Adam().LearningRate(0.001).Done()
optimizers.RMSProp().LearningRate(0.001).Done()
// 注意：不存在 optimizers.Momentum()

// 损失函数——包名是 losses，不是 loss
import "github.com/gomlx/gomlx/pkg/ml/train/losses"
losses.MeanSquaredError(labels, predictions)          // 参数是 []*Node 切片
losses.SparseCategoricalCrossEntropyLogits(labels, logits) // 参数是 []*Node 切片

张量（tensors）

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

t := tensors.FromValue(data)           // Go 值 → 张量（多维切片）
t := tensors.FromAnyValue(value)        // any 值 → 张量（含标量）
t.Value()                              // 张量 → any
t.FinalizeAll()                        // 释放设备内存
t.MaterializeLocal()                   // 搬回本地（无返回值）
t.Shape()                              // 获取形状
tensors.ToScalar[float64](t)           // 标量取值
tensors.MutableFlatData[float64](t, fn) // 原位修改（返回 error）

作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》

[上一章：第15章 GoMLX 的边界——WebAssembly、Gemma 推理与未来] | [全书完]