gomlx/gomlx

解决的问题

GoMLX 为 Go 编程语言提供了一个全面的机器学习和数学框架，作为 PyTorch、JAX 或 TensorFlow 等流行框架的 Go 原生替代方案。它使开发者能够在 Go 生态系统内训练、微调和执行机器学习模型，减少对 Python 进行 ML 开发和生产部署的依赖。

工作原理

GoMLX 作为一个中间 API，提供一系列 ML 层和工具库，并通过可互换的后端执行这些操作：

• XLA 后端：使用 OpenXLA 实现高性能 JIT 编译，支持 CPU、GPU（Nvidia、AMD ROCm、Intel、Mac）和 TPU。支持跨多个加速器的分布式执行。
• Go 后端：纯 Go 实现，高度可移植，支持 WASM 用于浏览器执行，并针对某些架构使用 SIMD 优化。
• ONNX 后端：利用 ONNX Runtime 执行计算，并支持将模型保存为 .onnx 格式。

该框架包含自动微分（autodiff）、变量管理（Store 和 Scope）系统，以及大量可微分操作符和 ML 层。

适用人群

专为希望在不离开 Go 生态系统的情况下构建、实验并部署机器学习模型的 Go 开发者设计，也适用于探索非传统优化器或正则化方法的研究人员。

主要亮点

• HuggingFace 集成：与 go-huggingface 兼容，支持下载模型、分词器、数据集，并从 GGUF 或 safetensors 读取参数。
• 多样化的层支持：包含 FFN、卷积、多头注意力、LSTM，以及 KAN（Kolmogorov-Arnold 网络）和 VNN（向量神经网络）等特殊层。
• 灵活的执行方式：通过 XLA 支持多种硬件目标，包括多 GPU/TPU 分布式执行。
• 开发者工具：包含用于检查检查点和生成图表的命令行工具，并通过 GoNB 与 JupyterLab 集成。
• 高级训练功能：支持梯度检查点，可在大规模模型训练中以内存换计算。