GoMLX，一个加速的机器学习和数学框架

📖 关于 GoMLX

GoMLX 是一套易于使用的机器学习和通用数学库及工具。 它可以被视为 Go 语言版的 PyTorch/Jax/TensorFlow。

它可用于训练、微调、修改和组合机器学习模型。 提供了使这些操作变得简单的所有工具：从完整的可微分算子集， 到在笔记本中进行训练时用于绘制指标的 UI 工具。

它几乎可以在任何运行 Go 的地方运行，使用纯 Go 后端实现。 甚至可以在浏览器中通过 WASM 运行（查看用 GoMLX 创建的演示）。 很可能它也能在嵌入式设备上运行（参见 Tamago）。

此外，它还支持基于 OpenXLA 的高度优化后端引擎， 该引擎利用即时编译技术，可在 CPU、GPU（Nvidia，以及可能的 AMD ROCm 和 Intel，还有 Mac 设备）以及 Google 的 TPU 上运行。 它还支持现代分布式执行功能（新功能，仍在积极改进中），可通过 XLA Shardy 实现多 TPU 或多 GPU 分布式训练， 这是对 GSPMD 分布式策略 的进一步发展。

该引擎与 Google 的 Jax、 TensorFlow 和 Pytorch/XLA 所使用的引擎相同， 在许多情况下性能也相当。 使用此后端可以训练大型模型或处理大规模数据集。

[!Tip]

• 请参阅我们的 🎓 教程 🎓
• 参阅 Eli Bendersky 的博客文章“GoMLX：无需 Python 的 Go 语言机器学习”（略显过时，但仍具参考价值）
• 一个针对 Kaggle 狗猫分类问题的引导式示例：狗猫分类。
• 一份简单的 GoMLX 幻灯片介绍，包含少量示例代码。

🗺️ 概述

GoMLX 是一个功能齐全的机器学习框架，支持从底层到高层的各种知名机器学习组件。 然而，它目前仍然只是主流机器学习库/框架（如 TensorFlow、Jax 或 PyTorch）所提供功能的一部分。

使用 GoMLX 开发的示例

• 🚀 新增 🚀 KaLM-Gema3 12B 参数：腾讯排名靠前的用于 RAG 的句子编码器，使用 go-huggingface 加载模型和分词器，并通过 GoMLX 执行。
• 🚀 新增 🚀 Gemma 3 270M：演示了使用 onnx-community/gemma-3-270m-it-ONNX 模型并通过 ONNX 转换后的文本生成（LLM），采用 GoMLX 运行。该示例使用 gomlx/onnx-gomlx 包进行模型转换，并借助 gomlx/go-huggingface 下载模型及运行分词器。
• 🚀 新增 🚀 GPT-2：展示了使用新的（实验性）变换器和生成器包进行文本生成。
• 🚀 新增 🚀 BERT-base-NER：一个针对命名实体识别任务微调的 BERT-base 模型。该模型同样是由 HuggingFace 上的 dslim/bert-base-NER 模型 转换为 ONNX 格式的。

• 🚀 新增 🚀 MixedBread Reranker v1：一个交叉编码器重排序示例，详情请参见 HuggingFace MixedBread Reranker v1 页面。该示例使用 gomlx/onnx-gomlx 包进行模型转换，并借助 gomlx/go-huggingface 下载模型及运行分词器。

• Adult/Census 模型;
• KAN 如何学习？;
• Cifar-10 示例;
• MNIST 示例（仅限库和命令行版本)
• 狗猫分类器示例;
• IMDB 电影评论示例;
• 牛津花卉 102 数据集的扩散模型（生成随机花卉）;
  • 流匹配研究笔记本，基于 Meta 的 "流匹配指南与代码"。
• OGBN-MAG 的 GNN 模型（实验性）。
• 最后，还有一个简单的 合成线性模型，供好奇者查看一个最基础的简单模型。
• 神经风格迁移十周年纪念：查看使用 GoMLX 编写的演示，基于 原始论文。
• 三元组损失：包含多种负采样策略以及不同的距离度量方法。
• AlphaZero AI 用于蜂巢棋游戏：它使用一个简单的 GNN 来评估棋盘上的局面。其中包含一个 WASM 演示（在浏览器中运行 GoMLX！) 和一个命令行界面，供玩家测试自己的技巧！

后端

GoMLX 是一个友好的“中间 ML API”，提供通用的 API 和一系列 ML 层等组件。但它本身并不执行任何计算：它依赖于不同的后端来编译并在各种硬件上执行计算。

目前有一个通用的后端接口（位于 github.com/gomlx/gomlx/backends，但很快将迁移到独立的仓库），并有三种不同的实现：

1. xla：适用于 CPU、GPU 和 TPU 的 OpenXLA 后端。就现有技术而言属于最先进的水平，但仅支持静态形状。 目前支持 linux/amd64、linux/arm64（CPU）以及 darwin/arm64（CPU）。使用 go-xla 提供的 Go 版本 API。
2. go：纯 Go 后端（无 C/C++ 依赖）：速度较慢但非常便携（可编译为 WASM、Windows 等）：
   • SIMD 支持正在推进中（参见 Go 的 SIMD 讨论 以及正在开发的 go-highway）；
   • 🚀 新增 🚀：新增对部分 融合操作 和某些量化类型的支持，从而在某些情况下显著提升性能。
   • 另请参阅 GoMLX 编译为 WASM 为蜂巢棋游戏的 AI 提供支持。
   • 计划支持动态形状（可能在 2026 年年中）。
3. 🚀 新增 🚀 go-darwinml：Apple CoreML 的 Go 绑定，支持 Metal 加速、MLX 以及所有 DarwinOS 相关的后端。

亮点

• 使用 onnx-gomlx 将 ONNX 模型转换为 GoMLX：既可以作为 onnxruntime 的替代方案（利用 XLA），也可以用于进一步微调模型。同时，还可以通过 go-huggingface 轻松从 HuggingFace 下载 ONNX 模型文件。
• Docker "gomlx_jupyterlab"，集成了 JupyterLab 和 GoNB（Jupyter 笔记本的 Go 内核）。
• 自动微分：目前仅支持梯度计算，暂不支持雅可比矩阵。
• 上下文管理：自动管理机器学习模型中的变量。
• 包含多种流行机器学习“层”的 ML 层库：前馈神经网络层、各类激活函数、层归一化和批归一化、卷积、池化、Dropout、多头注意力机制（用于 Transformer 层）、LSTM、KAN（B 样条、GR-KAN/KAT 网络、离散 KAN、分段线性 KAN）、分段线性层（用于校准和归一化）、各种正则化方法、可逆/可微分快速傅里叶变换、可学习有理函数（既可用于激活函数，也可用于 GR-KAN/KAT 网络）、VNN（向量神经网络）用于 SO(3) 等变/不变层等。
• 训练库，带有格式化的打印输出，并支持在 Jupyter 笔记本中使用 GoNB, a Go Kernel 绘制图表。
  • 此外，还提供多种调试工具：收集特定节点的值以供绘图、在训练过程中简单记录节点的值，以及生成创建节点的代码堆栈跟踪信息。
• gomlx_checkpoints 是一个命令行工具，用于检查训练模型的检查点，并生成图表，使用 Plotly 绘制损失曲线和任意评估指标。请参阅训练会话示例，其中展示了训练过程中学习率变化的效果。该工具还允许将不同模型的训练过程绘制在同一张图上，以便比较它们的演变。
• SGD 和 Adam（AdamW 和 Adamax）优化器。
• 多种损失函数和评估指标。
• 可直接使用的预训练模型：InceptionV3（图像模型），以及通过 onnx-gomlx 获取的更多来自 HuggingFace 的模型。同时，可以使用 go-huggingface 轻松从 HuggingFace 下载 ONNX 模型文件。
• 将 Numpy 数组读取为 GoMLX 张量——请参阅包 github.com/gomlx/gomlx/pkg/core/tensors/numpy。
• （实验性）支持 PJRT 的静态链接：虽然构建 Go 程序的速度较慢，但部署时无需在目标机器上安装 PJRT 插件。这需要您从 XLA 源码编译自己的静态 PJRT 插件。可以使用 go build --tags=pjrt_cpu_static 命令，或通过 import _ "github.com/gomlx/gomlx/backends/xla/cpu/static" 来启用。
• 自动安装 XLA PJRT 插件（适用于 CPU、GPU 和 TPU；支持 Linux 和 Mac），安装位置为用户的本地库目录（Linux 下为 $HOME/.local/lib，Mac 下为 $HOME/Library/Application Support/XLA）。可通过设置 GOMLX_NO_AUTO_INSTALL 或编程方式调用 xla.EnableAutoInstall(false) 来禁用此功能。
• 分布式执行（跨多个 GPU 或 TPU），用户只需提供少量提示即可。只需配置好分布式数据集，训练器便会自动完成其余工作。请参阅 UCI-Adult 示例 中的代码更改。此功能目前仍处于实验阶段，请报告任何问题并帮助我们改进它。

👥 支持

• 在 Slack 频道 #gomlx 进行讨论（您可以在此加入 Slack 服务器）。
• 问答与讨论
• 问题反馈
• 随机项目头脑风暴：只需发起一个问答，我很乐意在 Discord 或视频通话中与您交流。
• Google Groups: groups.google.com/g/gomlx-discuss

🛠️ + ⚙️ 安装

对于大多数用户来说，无需安装。

对于 XLA，系统默认会自动安装所需的 XLA PJRT 插件（适用于 CPU、GPU 和 TPU；支持 Linux 和 Mac），安装位置为用户的本地库目录（Linux 下为 $HOME/.local/lib/go-xla；Mac 下为 $HOME/Library/Application Support/go-xla；Windows 下为 $HOME\AppData\Local\go-xla）。可以通过设置 GOMLX_NO_AUTO_INSTALL 或编程方式调用 xla.EnableAutoInstall(false) 来禁用此功能。

如果您希望手动预安装以构建生产级 Docker 镜像、指定特定版本，或进行其他自定义设置，请参阅 github.com/gomlx/go-xla 以获取详细信息，其中包含易于理解的简单安装程序。

如果您只想使用纯 Go 后端，只需导入 _ "github.com/gomlx/gomlx/backends/simplego" 即可，无需安装任何额外组件。

🐳 预构建的 Docker

最简单的入门方式就是拉取这个包含 GoMLX + JupyterLab + GoNB（Jupyter 的 Go 内核）以及 Nvidia 的 CUDA 运行时（用于可选的 GPU 支持）的 Docker 镜像——镜像大小约为 5GB。CUDA 运行时已预先安装，无需额外操作。

在您希望在 Jupyter 中可见的目录下执行以下命令：

如果需要 GPU 支持，请在下面的 docker run 命令中添加 --gpus all 标志。

docker pull janpfeifer/gomlx_jupyterlab:latest
docker run -it --rm -p 8888:8888 -v "${PWD}":/home/jupyter/work janpfeifer/gomlx_jupyterlab:latest

终端中会显示一个以 127.0.0.1:8888 开头的 URL（包含所需的访问令牌），您可以在浏览器中打开该链接。

您还可以从那里打开并交互式地体验教程，该教程已包含在 Docker 的 Projects/gomlx/examples/tutorial 目录中。

更多关于该 Docker 的详细信息请参见 docker/jupyterlab/README.md。

该 Docker 也支持 Windows 系统：Docker Desktop 底层使用 WSL2 技术。

🧭 教程

请参阅此处的教程。它涵盖了各个方面的内容。

之后，可以查看位于examples/目录中的示例。

该库本身文档齐全（如果发现遗漏之处，请提交问题），代码也相对容易阅读。 Godoc 文档可在 pkg.go.dev 上找到。

最后，欢迎随时提问：在时间允许的情况下（非工作时间），我非常乐意提供帮助： 我经常在线于 Slack 频道 #gomlx； 或者访问 groups.google.com/g/gomlx-discuss。

推理与生产部署

目前，模型的推理或服务是通过使用创建模型时所用的 Go 代码，并结合包含训练权重和超参数的检查点来实现的。 换句话说，就是使用与训练相同的工具来进行推理。

例如，直接构建一个包含预训练模型的 Docker 容器并从中提供服务，或者将其集成到您自己的应用程序中，都是非常简单的操作。

关于如何实现这一点以及将模型推理导出为库的简单示例，请参阅 .../examples/cifar/classifer, 以及在Cifar-10 示例的最后几节中的使用方法。

未来，我们计划还将模型导出为 ONNX 或 XLA 的 StableHLO 格式，这样就可以直接使用支持这些格式的工具进行服务， 而无需链接 GoMLX——这将略微减小可执行文件的大小。

🎯 长期目标

1. 在 Go 中构建和训练模型——而非 Python（或其他语言）——重点在于：
   • 易于阅读和理解，使用户能够形成对当前操作过程正确且透明的心理模型。即使这意味着编写时会更冗长。
   • 清晰、可分离的 API：每个 API 应尽可能自成一体且相互解耦。
   • 可组合性：任何组件都应可替换，以便进行自定义和实验。 这意味着有时需要更多的编码工作（不存在一个能完成所有任务的“魔术”训练对象）， 但这样做可以使流程更加清晰，并便于用自定义版本替换部分组件。
   • 及时更新的文档：如果文档缺失或编写不佳，就如同代码不存在一样。
   • 清晰且可操作的错误报告
2. 成为一个高效的研究和教育平台，用于试验新的机器学习思想并进行学习。
   • 支持多设备上的镜像训练以及各种形式的分布式训练（模型并行和/或数据并行）， 特别是支持大型语言模型及类似大规模模型的训练。
3. 成为一个稳健可靠的生产级平台。一些子目标包括：
   • 支持现代加速硬件，如 TPU 和 GPU。
   • 提供除 XLA 之外的多种后端，例如：llamacpp、WebNN（通过 Wasm）、纯 Go 实现等。
   • 导入来自Hugging Face Hub的预训练模型并支持微调——ONNX 版本已经在 onnx-gomlx 中对许多模型生效。
   • 将模型编译为 C 库和/或 WebAssembly 二进制文件，以便在任何地方（任何语言）进行链接和推理使用。

🤔 常见问题解答

• GoMLX 使用哪些环境变量？
  • GOMLX_BACKEND：定义要使用的后端引擎（如果使用 backends.New()）。其值格式为“[:]”，其中配置部分是可选的。示例：
    • GOMLX_BACKEND=go：使用 SimpleGo 后端，即纯 Go 实现，具有很好的可移植性但速度较慢。
    • GOMLX_BACKEND="xla:cpu"：使用 XLA 后端（更快的后端，目前仅支持 Linux 系统）运行 CPU 计算。
    • GOMLX_BACKEND="xla:cuda"：使用 XLA 后端运行 Nvidia CUDA。
    • GOMLX_BACKEND="xla:/path/to/my/pjrt_plugin.so"：使用 XLA 并加载任意 PJRT 插件。PJRT 是 XLA 的插件系统，用于支持不同的硬件。 用户可以安装为 NVIDIA GPU 构建的 PJRT（有相应的安装脚本），也有适用于 ROCm 的 PJRT（作者未测试过）， 以及适用于 TPU（Google Cloud）的 PJRT，甚至还有针对新型加速器（例如：TensTorrent XLA）开发的 PJRT。
  • PJRT_PLUGIN_LIBRARY_PATH：底层 XLA 后端会将此变量视为额外的目录，用于搜索插件位置。 它会依次查找系统的库路径（$LD_LIBRARY_PATH、/etc/ld.so.conf）、默认路径 /usr/local/lib/gomlx/pjrt，以及用户设置的 $PJRT_PLUGIN_LIBRARY_PATH。
  • GOMLX_NO_AUTO_INSTALL：若设置为 1，GoMLX 在运行于尚未安装 PJRT 的系统上时，将不会自动安装 PJRT。
  • XLA_FLAGS：用于控制 XLA 后端的可选参数。应以分号（;）分隔的选项列表形式设置。若设置为 --help，后端会打印出所有选项的帮助信息。 更多信息请参阅XLA 标志指南页面。
• 使用 GoMLX 时应包含哪些后端？
  • 建议使用 import _ "github.com/gomlx/gomlx/backends/default"，这将导入 xla（或别名 stablehlo）和 go（SimpleGo）后端。 如果在编译时添加 -tags=noxla 参数，则不会包含 XLA 后端。
  • 使用 import _ "github.com/gomlx/gomlx/backends/simplego" 仅导入 go 后端（无 C++ 依赖）。
  • 使用 import _ "github.com/gomlx/gomlx/backends/xla" 仅导入 XLA 后端。
• AI 上下文文件（Gemini/Claude Code/Conductor 等）存放在哪里？
  • 为避免项目根目录过于杂乱，这些文件已被移至 .agents/ 目录，并已加入 .gitignore 文件中。 这样一来，用户只需将所需的 AI 配置文件以符号链接的形式指向本地副本的根目录，即可轻松使用它们。

🤝 合作

该项目欢迎所有感兴趣的人士贡献代码。 许多部分尚未最终确定，因此对于那些具备 Go、机器学习和 API 设计经验的人来说，仍有大量改进和重新设计的空间。 请参阅TODO 文件获取灵感。

目前尚未制定治理规范。

如需联系，请参阅上方的“支持”部分！（可通过 Slack 频道或 Google Groups）

💖 支持本项目

如果您觉得本项目有所帮助，请考虑通过GitHub Sponsors支持我们的工作。

您的捐助将帮助我们（目前主要是我）投入更多时间进行维护， 并为整个 GoMLX 生态系统添加新功能。

同时，这笔资金也将用于获取更多硬件资源的使用权（购买或云服务），以提升跨平台兼容性，例如：ROCm、Apple Metal（GPU）、多 GPU/TPU、NVIDIA DGX Spark、Tenstorrent 等。

🚀 高级主题

• 更新日志
• 待办事项
• 错误处理
• 开发指南

💖 感谢

• Go
• OpenXLA
• TensorFlow
• Jax
• PyTorch
• ONNX
• HuggingFace

⚖️ 许可证

版权归2025年扬·普费弗所有

GoMLX 根据 Apache许可证第2.0版 进行分发。 除非另有明确说明，否则任何有意提交以纳入本项目的贡献，都将依据 Apache许可证第2.0版 进行许可，且不附加任何额外条款或条件。

GoMLX 快速上手指南

GoMLX 是专为 Go 语言设计的机器学习与数学框架，被誉为"Go 版本的 PyTorch/JAX/TensorFlow"。它支持纯 Go 后端（可运行于浏览器/WASM）以及基于 OpenXLA 的高性能后端（支持 CPU、GPU 和 TPU 加速）。

环境准备

系统要求

• 操作系统：Linux (amd64/arm64), macOS (arm64/amd64), Windows (amd64)
• Go 版本：建议使用 Go 1.21 或更高版本
• 硬件加速（可选）：
  • 若使用 GPU/TPU 加速，需安装对应的 NVIDIA CUDA 驱动或 Google TPU 环境。
  • 若仅使用纯 Go 后端，无需额外硬件依赖，甚至可在浏览器中通过 WASM 运行。

前置依赖

• 确保已安装 Git。
• 若计划使用 Jupyter Notebook 进行交互式开发，推荐安装 GoNB。

安装步骤

1. 初始化 Go 模块

在你的项目目录中初始化 Go 模块（如果尚未初始化）：

go mod init my-gomlx-project

2. 安装 GoMLX 核心库

使用 go get 命令安装 GoMLX：

go get github.com/gomlx/gomlx

3. 安装后端引擎（按需选择）

方案 A：纯 Go 后端（默认，高兼容性）

无需额外操作，代码默认使用纯 Go 后端，适合开发调试、WASM 部署或无 GPU 环境。

方案 B：OpenXLA 后端（高性能，支持 GPU/TPU）

若需利用 GPU 或 TPU 加速训练大模型，需引入 XLA 后端包：

go get github.com/gomlx/gomlx/backends/xla

注意：使用 XLA 后端时，请确保系统已安装相应的 CUDA 工具包（针对 NVIDIA GPU）或 TPU 运行时库。

方案 C：ONNX 模型支持（可选）

若需加载 HuggingFace 上的 ONNX 模型：

go get github.com/gomlx/onnx-gomlx
go get github.com/gomlx/go-huggingface

基本使用

以下是一个最简单的示例，展示如何创建一个线性模型并执行前向传播计算。

示例代码：构建简单的线性层

创建文件 main.go，写入以下内容：

package main

import (
	"fmt"
	"github.com/gomlx/gomlx/backends"
	"github.com/gomlx/gomlx/graph"
	"github.com/gomlx/gomlx/layers"
	"github.com/gomlx/gomlx/types/tensors"
)

func main() {
	// 1. 获取默认后端 (默认为纯 Go 后端，若安装了 xla 包且环境支持可自动切换)
	backend := backends.DefaultBackend()

	// 2. 创建计算图上下文
	ctx := graph.NewContext(backend)

	// 3. 定义输入张量 (形状: [batch_size=2, features=3])
	input := tensors.FromScalarSlice(backend, []float32{1, 2, 3, 4, 5, 6}, []int{2, 3})

	// 4. 在计算图中构建模型：一个简单的全连接层 (Dense Layer)
	// 输入维度 3 -> 输出维度 1
	output := layers.Dense(ctx, input, 1, nil, "linear_layer")

	// 5. 执行计算并获取结果
	result := output.Value()
	
	fmt.Printf("输入形状：%v\n", input.Shape())
	fmt.Printf("输出结果：%v\n", result.Slice())
	
	// 清理资源
	ctx.Done()
}

运行程序

go run main.go

下一步建议

• 教程学习：访问官方 Tutorial 了解完整的训练流程。
• 实战案例：参考 examples 目录下的 MNIST 或 GPT-2 示例。
• 交互式开发：使用 Docker 启动集成环境：

  docker run -p 8888:8888 janpfeifer/gomlx_jupyterlab