GitHub - fumi-engineer/machine_learning

┌─────────────────────────────────────────────────────────────┐
│                    MoE Transformer                          │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │
│  │  Embedding  │→ │ Blocks × 30 │→ │  LM Head    │         │
│  └─────────────┘  └─────────────┘  └─────────────┘         │
│                          ↓                                  │
│  ┌──────────────────────────────────────────────────────┐  │
│  │              Transformer Block                        │  │
│  │  ┌─────────┐    ┌───────────┐    ┌─────────┐         │  │
│  │  │ RMSNorm │ →  │    MQA    │ →  │ RMSNorm │ → MoE   │  │
│  │  └─────────┘    │ (12Q/1KV) │    └─────────┘         │  │
│  │       ↑         └───────────┘         ↑     ↓        │  │
│  │       └──────── residual ─────────────┴─ residual ───│  │
│  └──────────────────────────────────────────────────────┘  │
│                          ↓                                  │
│  ┌──────────────────────────────────────────────────────┐  │
│  │                   MoE Layer                           │  │
│  │  ┌────────┐   ┌─────────────────────────────────┐    │  │
│  │  │ Router │ → │ Expert 0  │ ... │  Expert 15   │    │  │
│  │  │ top-4  │   │ (SwiGLU)  │     │  (SwiGLU)    │    │  │
│  │  └────────┘   └─────────────────────────────────┘    │  │
│  └──────────────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────────────┘

machine_learning/
├── rust/             # Rust implementation
│   ├── nn-core/      # Model, tensor ops, training
│   ├── nn-cuda/      # CUDA FFI bindings
│   └── nn-ffi/       # FFI bridge + GpuTrainer
├── go/               # Go implementation
│   ├── tensor/       # Tensor operations
│   ├── cuda/         # cgo CUDA bindings
│   ├── layer/        # Neural network layers
│   ├── model/        # MoE Transformer model
│   └── train/        # Training pipeline
├── python/           # Python implementation
│   ├── nn/           # Neural network module
│   ├── cuda/         # ctypes CUDA bindings
│   └── tests/        # pytest tests
├── cuda/             # Shared CUDA kernels (.cu, stub.c)
├── benchmarks/       # Cross-language benchmarks
│   ├── rust/         # Criterion benchmarks
│   ├── go/           # testing.B benchmarks
│   └── python/       # timeit + NumPy benchmarks
├── docs-jp/          # 日本語ドキュメント
├── docs-en/          # English documentation
└── Cargo.toml        # Rust workspace

# Build
cargo build --release

# Test (53 tests)
cargo test

# Clippy
cargo clippy --all-targets

cd go

# Test
go test ./...

# Build (requires CUDA library)
go build ./...

cd python

# Install
pip install -e ".[dev]"

# Test (42 tests)
pytest

# Force CPU-only (Rust)
CUDA_PATH="" cargo build --release

// Create tiny model for testing
let model = MoETransformer::tiny();

// Forward pass
let token_ids = vec![1, 2, 3, 4];
let logits = model.forward_ids(&token_ids, 1, 4);
// → [1, 4, vocab_size] logits

// Create tiny model for testing
model := model.NewTiny()

// Forward pass
tokenIDs := []int{1, 2, 3, 4}
logits := model.ForwardIDs(tokenIDs, 1, 4)
// → [1, 4, vocab_size] logits

from nn import MoETransformer

# Create tiny model for testing
model = MoETransformer.tiny()

# Forward pass
token_ids = [1, 2, 3, 4]
logits = model.forward_ids(token_ids, batch=1, seq_len=4)
# → [1, 4, vocab_size] logits

# Run benchmarks
cd benchmarks
./run_all.sh              # All languages
./run_all.sh --rust-only  # Rust only
./run_all.sh --go-only    # Go only
./run_all.sh --python-only # Python only