A deep learning training framework built with MoonBit, featuring tape-based autograd and PyTorch-like API (CPU backend)
Dependencies
moon add chnlkw/moonxi-netNote: This package works with the standard MoonBit toolchain — no CUDA or GPU required.
///|
fn[T : @tensor.Tensor + @tensor.BlasTensor] train_linear() -> Array[T] {
let w = @grad.no_grad(T::zeros([1, 1]))
let b = @grad.no_grad(T::zeros([1, 1]))
let x = @grad.no_grad(
T::from_host(FixedArray::makei(5, i => Float::from_int(i + 1)), [5, 1]),
)
let y = @grad.no_grad(
T::from_host(FixedArray::makei(5, i => Float::from_int((i + 1) * 3 + 1)), [
5, 1,
]),
)
let params : Array[@grad.Grad[T]] = [w, b]
for _epoch in 0500 {
@grad.clear_tape()
for p in params { p.grad = Some(None) }
let loss = x.matmul(w).add(b).sub(y).square().mean()
loss.backward()
for p in params {
match p.grad {
Some(Some(g)) => p.value = p.value.sub(g.scale(0.001))
_ => ()
}
}
}
params.map(g => g.value)
}moon run moonxi-net/examples/linear --target native
# Learned: w=3.00, b=1.00{
"import": [
"chnlkw/moonxi-net" @tensor,
"chnlkw/moonxi-net/nparray" @nparray,
"chnlkw/moonxi-net/grad" @grad,
"chnlkw/moonxi-net/model" @model,
"chnlkw/moonxi-net/optimizer" @optimizer,
"chnlkw/moonxi-net/loss" @loss,
"chnlkw/moonxi-net/train" @train,
"chnlkw/moonxi-net/dataloader" @dl
]
}| Package | Alias | Description |
|---|---|---|
| moonxi-net | @tensor | Core traits: Tensor, BlasTensor, ImageTensor, ImageBackwardOps |
| moonxi-net/nparray | @nparray | CPU tensor backend (NpArray implements all tensor traits) |
| moonxi-net/grad | @grad | Tape-based autograd engine (Grad[T], backward(), clear_tape()) |
| moonxi-net/model | @model | Neural network layers: Linear[T], Conv2d[T], ResNet18[T], MLP[T] |
| moonxi-net/optimizer | @optimizer | Optimizers: Momentum SGD, Adam, RMSprop (with gradient clipping) |
| moonxi-net/loss | @loss | Loss functions: cross-entropy, MSE |
| moonxi-net/train | @train | Training loop utilities, Experiment config, run_cpu |
| moonxi-net/dataloader | @dl | DataLoader with Fisher-Yates shuffle and mini-batch iteration |
| moonxi-net/datasets/mnist | — | MNIST dataset loader |
| moonxi-net/datasets/cifar10 | — | CIFAR-10 dataset loader |
///|
fn main {
// Create tensors using NpArray (CPU backend)
let w = @grad.no_grad(@nparray.NpArray::zeros([10, 10]))
let b = @grad.no_grad(@nparray.NpArray::zeros([10]))
// Build computation graph, compute loss, backprop
@grad.clear_tape()
let loss = /* ... your forward pass ... */
loss.backward()
// Access gradients
match w.grad {
Some(Some(g)) => w.value = w.value.sub(g.scale(0.01))
_ => ()
}
}# From this directory
moon test --target native
# Build
moon build --target nativepub(open) trait BlasTensor {
fn matmul(Self, Self) -> Self
fn transpose(Self) -> Self
}pub(open) trait ImageBackwardOps {
fn conv2d_backward_data(grad_output : Self, weight : Self, input : Self, stride : Int, padding : Int) -> Self
fn conv2d_backward_weight(grad_output : Self, input : Self, weight : Self, stride : Int, padding : Int) -> Self
fn conv2d_backward_bias(grad_output : Self) -> Self
fn relu_backward(grad_output : Self, input : Self) -> Self
fn batchnorm_backward(grad_output : Self, input : Self, gamma : Self, save_mean : Self, save_inv_var : Self, eps : Float) -> (Self, Self, Self)
fn maxpool2d_backward(grad_output : Self, input : Self, kernel_size : Int, stride : Int) -> Self
fn adaptive_avg_pool2d_backward(grad_output : Self, input : Self) -> Self
fn softmax_ce_backward(logits : Self, targets : Self) -> Self
fn softmax_ce_backward_labels(logits : Self, labels : Self, num_classes : Int) -> Self
}pub(open) trait ImageTensor {
fn conv2d(Self, weight : Self, bias : Self, stride : Int, padding : Int) -> Self
fn relu(Self) -> Self
fn maxpool2d(Self, kernel_size : Int, stride : Int) -> Self
fn adaptive_avg_pool2d(Self, output_size : Int) -> Self
fn batchnorm_training(Self, gamma : Self, beta : Self, running_mean : Self, running_var : Self, momentum : Float, eps : Float) -> (Self, Self, Self)
fn batchnorm_inference(Self, gamma : Self, beta : Self, running_mean : Self, running_var : Self, eps : Float) -> Self
fn softmax_cross_entropy(Self, targets : Self) -> Self
fn cross_entropy_with_labels(Self, labels : Self, num_classes : Int) -> Self
}pub(open) trait Tensor {
fn dims(Self) -> FixedArray[Int]
fn zeros(dims : FixedArray[Int]) -> Self
fn zeros_like(Self) -> Self
fn from_host(data : FixedArray[Float], shape : Array[Int]) -> Self
fn square(Self) -> Self
fn sqrt(Self) -> Self
fn mean(Self) -> Self
fn scale(Self, Float) -> Self
fn mul_elem(Self, Self) -> Self
fn div_elem(Self, Self) -> Self
fn scalar(Float) -> Self
fn size(Self) -> Int
fn reduce_sum_to(Self, target_dims : FixedArray[Int]) -> Self
fn value(Self) -> TensorData
fn view(Self, new_shape : FixedArray[Int]) -> Self
fn add(Self, Self) -> Self
fn add_into(Self, Self) -> Unit
fn sub(Self, Self) -> Self
fn broadcast_to(Self, target_shape : FixedArray[Int]) -> Self
}impl BlasTensor for ShapeTensorimpl ImageBackwardOps for ShapeTensorfn batchnorm_backward(_grad_output : ShapeTensor, input : ShapeTensor, gamma : ShapeTensor, _save_mean : ShapeTensor, _save_inv_var : ShapeTensor, _eps : Float) -> (ShapeTensor, ShapeTensor, ShapeTensor)fn conv2d_backward_data(_grad_output : ShapeTensor, _weight : ShapeTensor, input : ShapeTensor, _stride : Int, _padding : Int) -> ShapeTensorfn conv2d_backward_weight(_grad_output : ShapeTensor, _input : ShapeTensor, weight : ShapeTensor, _stride : Int, _padding : Int) -> ShapeTensorfn maxpool2d_backward(_grad_output : ShapeTensor, input : ShapeTensor, _kernel_size : Int, _stride : Int) -> ShapeTensorfn softmax_ce_backward_labels(logits : ShapeTensor, _labels : ShapeTensor, _num_classes : Int) -> ShapeTensorimpl ImageTensor for ShapeTensorfn batchnorm_inference(input : ShapeTensor, _gamma : ShapeTensor, _beta : ShapeTensor, _running_mean : ShapeTensor, _running_var : ShapeTensor, _eps : Float) -> ShapeTensorfn batchnorm_training(input : ShapeTensor, _gamma : ShapeTensor, _beta : ShapeTensor, _running_mean : ShapeTensor, _running_var : ShapeTensor, _momentum : Float, _eps : Float) -> (ShapeTensor, ShapeTensor, ShapeTensor)fn conv2d(input : ShapeTensor, weight : ShapeTensor, bias : ShapeTensor, stride : Int, padding : Int) -> ShapeTensorfn cross_entropy_with_labels(_logits : ShapeTensor, _labels : ShapeTensor, _num_classes : Int) -> ShapeTensorimpl Tensor for ShapeTensorfn ShapeTensor::pool2d(self : ShapeTensor, kernel_size : Int, stride : Int, padding : Int) -> ShapeTensorA deep learning training framework built with MoonBit, featuring tape-based autograd and PyTorch-like API (CPU backend)
Dependencies