GoMLX：Go语言的机器学习和数学计算框架

// 计算图示例
func buildModel(ctx context.Context, graph *graph.Graph, inputs []*graph.Node) []*graph.Node {
    // 第一个隐藏层
    hidden1 := layers.FC(graph, inputs[0], 128, "hidden1")
    hidden1 = graph.ReLU(hidden1)
    
    // 第二个隐藏层
    hidden2 := layers.FC(graph, hidden1, 64, "hidden2")
    hidden2 = graph.ReLU(hidden2)
    
    // 输出层
    logits := layers.FC(graph, hidden2, 10, "logits")
    
    return []*graph.Node{logits}
}

// 自动微分示例
func trainStep(ctx context.Context, model *layers.Model, optimizer train.Optimizer, 
               inputs, labels []*graph.Node, trainMetrics *train.Metrics) {
    // 前向传播
    logits := model.Fwd(ctx, inputs...)
    
    // 计算损失
    loss := losses.CrossEntropyLogits(logits, labels)
    
    // 计算梯度
    gradients := graph.Gradient(ctx, loss, model.Variables()...)
    
    // 更新变量
    optimizer.UpdateVariables(gradients...)
    
    // 更新指标
    trainMetrics.Update(loss, logits, labels)
}

// 启用StableHLO后端
import _ "github.com/gomlx/gomlx/backends/stablehlo"

// 或者使用标签
// go run -tags=stablehlo main.go

package main

import (
    "context"
    "fmt"
    
    "github.com/gomlx/gomlx/graph"
    "github.com/gomlx/gomlx/layers"
    "github.com/gomlx/gomlx/ml/train"
    "github.com/gomlx/gomlx/ml/train/optimizers"
    "github.com/gomlx/gomlx/types/tensors"
)

func main() {
    // 创建后端和上下文
    backend := backends.New()
    ctx := context.Background()
    
    // 创建模型
    model := layers.Sequential(
        layers.FC(nil, 128, "fc1"),
        layers.ReLU(nil),
        layers.FC(nil, 64, "fc2"),
        layers.ReLU(nil),
        layers.FC(nil, 10, "fc3"),
    )
    
    // 创建优化器
    optimizer := optimizers.Adam().Build()
    
    // 训练循环
    for epoch := 0; epoch < 10; epoch++ {
        // 这里应该是数据加载和训练步骤
        // 简化的训练步骤示例
        trainStep(ctx, model, optimizer, inputs, labels)
    }
    
    fmt.Println("训练完成")
}

func trainStep(ctx context.Context, model *layers.Model, optimizer train.Optimizer, 
               inputs, labels []*graph.Node) {
    // 前向传播
    logits := model.Fwd(ctx, inputs...)
    
    // 计算损失
    loss := losses.CrossEntropyLogits(logits, labels)
    
    // 计算梯度
    gradients := graph.Gradient(ctx, loss, model.Variables()...)
    
    // 更新变量
    optimizer.UpdateVariables(gradients...)
}

package main

import (
    "context"
    "fmt"
    
    "github.com/gomlx/gomlx/backends"
    "github.com/gomlx/gomlx/graph"
    "github.com/gomlx/gomlx/layers"
    "github.com/gomlx/gomlx/ml/train"
    "github.com/gomlx/gomlx/ml/train/optimizers"
    "github.com/gomlx/gomlx/types/tensors"
)

// ConvNet 定义一个简单的卷积神经网络
func ConvNet(g *graph.Graph, inputs []*graph.Node) []*graph.Node {
    // 第一个卷积层
    conv1 := layers.Conv2D(g, inputs[0], 32, []int{3, 3}, []int{1, 1}, "SAME", "conv1")
    conv1 = g.ReLU(conv1)
    pool1 := g.MaxPool(conv1, []int{1, 2, 2, 1}, []int{1, 2, 2, 1}, "SAME")
    
    // 第二个卷积层
    conv2 := layers.Conv2D(g, pool1, 64, []int{3, 3}, []int{1, 1}, "SAME", "conv2")
    conv2 = g.ReLU(conv2)
    pool2 := g.MaxPool(conv2, []int{1, 2, 2, 1}, []int{1, 2, 2, 1}, "SAME")
    
    // 展平
    shape := g.Shape(pool2)
    batchSize := shape.Dim(0)
    flattened := g.Reshape(pool2, []int{batchSize, -1})
    
    // 全连接层
    fc1 := layers.FC(g, flattened, 128, "fc1")
    fc1 = g.ReLU(fc1)
    
    // 输出层
    logits := layers.FC(g, fc1, 10, "logits")
    
    return []*graph.Node{logits}
}

func main() {
    // 创建后端
    backend := backends.New()
    ctx := context.Background()
    
    // 创建模型
    model := layers.NewModel(ConvNet)
    
    // 创建优化器
    optimizer := optimizers.Adam().Build()
    
    // 训练模型（这里省略了数据加载和训练循环的具体实现）
    fmt.Println("卷积神经网络模型已创建")
}

# 创建新项目
mkdir my-gomlx-project
cd my-gomlx-project
go mod init example.com/my-gomlx-project

# 添加GoMLX依赖
go get github.com/gomlx/gomlx

# 如果需要高性能后端
go get github.com/gomlx/gomlx/backends/xla

# 或者使用StableHLO后端（测试版）
go get github.com/gomlx/gomlx/backends/stablehlo

package main

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

func main() {
    // 1. 创建后端
    backend := backends.New()
    ctx := context.Background()
    
    // 2. 创建计算图
    g := graph.New(ctx, backend)
    
    // 3. 创建输入张量
    x := g.Parameter(tensors.FromScalar(2.0))
    y := g.Parameter(tensors.FromScalar(3.0))
    
    // 4. 执行计算
    z := g.Add(x, g.Mul(x, y))  // z = x + x * y
    
    // 5. 执行计算图
    result := g.Run(z)
    
    // 6. 输出结果
    fmt.Printf("结果: %f\n", result.Value().(float64))  // 应该输出 8.0
}