cargo install axonc

git clone https://github.com/axon-lang/axon.git
cd axon
cargo build --release
# Binary is at target/release/axonc

axonc --version
# axonc 0.1.0

fn main() {
    println("Hello, Axon!");
}

axonc build hello.axon -o hello
./hello
# Hello, Axon!

axonc repl
>>> println("Hello from the REPL!")
Hello from the REPL!

axonc pkg new my_project
cd my_project

my_project/
├── Axon.toml          # Project manifest
├── src/
│   └── main.axon      # Entry point
└── tests/
    └── test_main.axon  # Test file

[package]
name = "my_project"
version = "0.1.0"
edition = "2026"

[dependencies]

fn main() {
    println("Hello from my_project!");
}

axonc pkg build
axonc pkg run
# Hello from my_project!

# Parse and check for errors
axonc check hello.axon

# Build an optimized binary
axonc build hello.axon -O 3 -o hello

# Emit LLVM IR for inspection
axonc build hello.axon --emit-llvm

axonc pkg build        # Build the project
axonc pkg run          # Build and run
axonc pkg test         # Run tests
axonc pkg fmt          # Format all source files
axonc pkg lint         # Lint all source files

cd editors/vscode
npm install
npm run build

axonc lsp

val x = 42;              // immutable, type inferred as Int32
val y: Float64 = 3.14;   // explicit type annotation
var counter = 0;          // mutable variable
counter += 1;

val name = "Axon";            // String
val active = true;            // Bool
val scores = [95, 87, 92];    // Vec<Int32>

fn add(a: Int32, b: Int32): Int32 {
    a + b   // last expression is the return value
}

fn greet(name: String) {
    println("Hello, {}!", name);
}

fn main() {
    val sum = add(3, 4);
    greet("World");
}

unsafe fn raw_pointer_access(ptr: *mut Float32): Float32 {
    // low-level memory access
}

val dec = 42;          // decimal
val hex = 0xFF;        // hexadecimal
val bin = 0b1010;      // binary
val oct = 0o77;        // octal
val sci = 1.5e10;      // scientific notation

val max = if a > b { a } else { b };

if score >= 90 {
    println("Excellent");
} else if score >= 70 {
    println("Good");
} else {
    println("Keep trying");
}

var i = 0;
while i < 10 {
    println("{}", i);
    i += 1;
}

for item in collection {
    println("{}", item);
}

for i in 0..10 {
    println("{}", i);
}

match value {
    0 => println("zero"),
    1 => println("one"),
    n => println("other: {}", n),
}

match option_val {
    Some(x) => println("Got {}", x),
    None => println("Nothing"),
}

model Point {
    x: Float64,
    y: Float64,
}

val p = Point { x: 1.0, y: 2.0 };
println("({}, {})", p.x, p.y);

extend Point {
    fn distance(&self, other: &Point): Float64 {
        val dx = self.x - other.x;
        val dy = self.y - other.y;
        (dx * dx + dy * dy).sqrt()
    }

    fn origin(): Point {
        Point { x: 0.0, y: 0.0 }
    }
}

enum Shape {
    Circle(Float64),
    Rectangle(Float64, Float64),
    Triangle { base: Float64, height: Float64 },
}

fn area(shape: Shape): Float64 {
    match shape {
        Shape.Circle(r) => 3.14159 * r * r,
        Shape.Rectangle(w, h) => w * h,
        Shape.Triangle { base, height } => 0.5 * base * height,
    }
}

trait Printable {
    fn to_string(&self): String;
}

extend Printable for Point {
    fn to_string(&self): String {
        format("({}, {})", self.x, self.y)
    }
}

fn print_all<T: Printable>(items: Vec<T>) {
    for item in items {
        println("{}", item.to_string());
    }
}

trait Drawable: Printable {
    fn draw(&self);
}

fn max<T: Ord>(a: T, b: T): T {
    if a > b { a } else { b }
}

model Pair<A, B> {
    first: A,
    second: B,
}

extend<A: Display, B: Display> Pair<A, B> {
    fn show(&amp;self) {
        println("({}, {})", self.first, self.second);
    }
}

// Tensor with known shape
val weights: Tensor<Float32, [784, 256]> = randn([784, 256]);

// Dynamic batch dimension with ?
val input: Tensor<Float32, [?, 784]> = load_batch();

// Matrix multiply — shapes checked at compile time
val output = input @ weights;   // Tensor<Float32, [?, 256]>

val a: Tensor<Float32, [3, 4]> = randn([3, 4]);
val b: Tensor<Float32, [5, 6]> = randn([5, 6]);
val c = a @ b;   // ERROR[E3001]: shape mismatch — inner dims 4 ≠ 5

fn find(haystack: Vec<Int32>, needle: Int32): Option<Int32> {
    for i in 0..haystack.len() {
        if haystack[i] == needle {
            return Some(i);
        }
    }
    None
}

fn read_config(path: String): Result<Config, IOError> {
    val file = File.open(path)?;    // propagate error with ?
    val data = file.read_all()?;
    parse_config(data)
}

mod math {
    pub fn square(x: Float64): Float64 {
        x * x
    }

    fn internal_helper() {
        // private — not visible outside this module
    }
}

use math.square;

fn main() {
    println("{}", square(4.0));   // 16.0
}

// Static shape — all dimensions known at compile time
val weights: Tensor<Float32, [784, 256]> = randn([784, 256]);

// Dynamic batch dimension (?)
val input: Tensor<Float32, [?, 784]> = load_batch();

// Fully dynamic shape
val dynamic: Tensor<Float32, [?, ?]> = some_function();

// Zeros and ones
val z = zeros([3, 4]);           // Tensor<Float32, [3, 4]>
val o = ones([256]);             // Tensor<Float32, [256]>

// Random initialization
val r = randn([128, 64]);       // normal distribution
val u = rand([10, 10]);         // uniform [0, 1)

// From data
val t = Tensor.from_vec([1.0, 2.0, 3.0, 4.0], [2, 2]);

// Range
val seq = arange(0, 10);        // [0, 1, 2, ..., 9]

// Identity matrix
val eye = Tensor.eye(4);       // 4×4 identity

// From file
val data = load_data("weights.npy");

val f16: Tensor<Float16, [1024]> = zeros([1024]);    // half precision
val f32: Tensor<Float32, [1024]> = zeros([1024]);    // single precision
val f64: Tensor<Float64, [1024]> = zeros([1024]);    // double precision
val i32: Tensor<Int32, [10]> = arange(0, 10);        // integer tensor

val a: Tensor<Float32, [2, 6]> = randn([2, 6]);
val b = a.reshape([3, 4]);      // Tensor<Float32, [3, 4]>
val c = a.reshape([12]);        // Tensor<Float32, [12]>
// val d = a.reshape([5, 5]);   // ERROR[E3002]: cannot reshape [2,6] (12 elements) to [5,5] (25 elements)

val m: Tensor<Float32, [3, 4]> = randn([3, 4]);
val mt = m.transpose();         // Tensor<Float32, [4, 3]>

// For higher-rank tensors, specify axes
val t: Tensor<Float32, [2, 3, 4]> = randn([2, 3, 4]);
val tp = t.permute([0, 2, 1]);  // Tensor<Float32, [2, 4, 3]>

val a: Tensor<Float32, [1, 3, 1, 4]> = randn([1, 3, 1, 4]);
val b = a.squeeze();            // Tensor<Float32, [3, 4]>

val c: Tensor<Float32, [3, 4]> = randn([3, 4]);
val d = c.unsqueeze(0);         // Tensor<Float32, [1, 3, 4]>

val a: Tensor<Float32, [2, 3]> = randn([2, 3]);
val b: Tensor<Float32, [2, 3]> = randn([2, 3]);

val cat = Tensor.cat([a, b], 0);    // Tensor<Float32, [4, 3]>
val stk = Tensor.stack([a, b], 0);  // Tensor<Float32, [2, 2, 3]>

val t: Tensor<Float32, [10, 20]> = randn([10, 20]);
val row = t[0];                  // Tensor<Float32, [20]>
val sub = t[2..5];               // Tensor<Float32, [3, 20]>

val a = randn([3, 4]);
val b = randn([3, 4]);

val sum  = a + b;     // element-wise addition
val diff = a - b;     // element-wise subtraction
val prod = a * b;     // element-wise multiplication (Hadamard)
val quot = a / b;     // element-wise division

// Scalar broadcasting
val scaled = a * 2.0;
val shifted = a + 1.0;

val x = randn([100]);

val s  = x.sin();
val c  = x.cos();
val e  = x.exp();
val l  = x.log();
val sq = x.sqrt();
val ab = x.abs();
val cl = x.clamp(-1.0, 1.0);

val h = relu(x);
val g = gelu(x);
val s = sigmoid(x);
val t = tanh(x);
val p = softmax(logits, dim: 1);

val t: Tensor<Float32, [4, 5]> = randn([4, 5]);

val total = t.sum();              // scalar
val row_sum = t.sum(dim: 1);     // Tensor<Float32, [4]>
val col_mean = t.mean(dim: 0);   // Tensor<Float32, [5]>
val max_val = t.max();            // scalar
val min_idx = t.argmin(dim: 1);  // Tensor<Int64, [4]>

val A: Tensor<Float32, [3, 4]> = randn([3, 4]);
val B: Tensor<Float32, [4, 5]> = randn([4, 5]);
val C = A @ B;    // Tensor<Float32, [3, 5]>

// Inner dimensions must match
val D: Tensor<Float32, [4, 6]> = randn([4, 6]);
// val E = A @ D;   // ERROR[E3001]: matmul shape mismatch — inner dims 4 ≠ 4... wait
//                   // actually [3,4] @ [4,6] works. Let's show a real error:
val F: Tensor<Float32, [5, 6]> = randn([5, 6]);
// val G = A @ F;   // ERROR[E3001]: matmul requires inner dims to match: 4 ≠ 5

val batch_a: Tensor<Float32, [?, 8, 64]> = randn([32, 8, 64]);
val batch_b: Tensor<Float32, [?, 64, 32]> = randn([32, 64, 32]);
val batch_c = batch_a @ batch_b;   // Tensor<Float32, [?, 8, 32]>

val M = randn([4, 4]);

val d = M.det();              // determinant
val inv = M.inv();            // inverse
val (Q, R) = M.qr();         // QR decomposition
val (U, S, V) = M.svd();     // singular value decomposition
val eig = M.eigenvalues();   // eigenvalues
val tr = M.trace();           // trace
val dp = a.dot(b);            // dot product (1D tensors)

val cpu_tensor = randn([1024, 1024]);

// Move to GPU
val gpu_tensor = cpu_tensor.to_gpu();

// Compute on GPU
val result = gpu_tensor @ gpu_tensor;

// Move back to CPU for I/O
val cpu_result = result.to_cpu();
println("{}", cpu_result);

// ✓ Shapes match
val a: Tensor<Float32, [3, 4]> = randn([3, 4]);
val b: Tensor<Float32, [4, 5]> = randn([4, 5]);
val c = a @ b;    // OK: [3,4] @ [4,5] → [3,5]

// ✗ Shape mismatch
val d: Tensor<Float32, [3, 4]> = randn([3, 4]);
val e: Tensor<Float32, [5, 6]> = randn([5, 6]);
// val f = d @ e;   // ERROR[E3001]: matmul inner dim mismatch: 4 ≠ 5

// ✗ Invalid reshape
val g: Tensor<Float32, [2, 3]> = randn([2, 3]);
// val h = g.reshape([2, 2]);  // ERROR[E3002]: element count mismatch: 6 ≠ 4

// ✗ Element-wise shape mismatch
val i: Tensor<Float32, [3, 4]> = randn([3, 4]);
val j: Tensor<Float32, [3, 5]> = randn([3, 5]);
// val k = i + j;   // ERROR[E3003]: broadcast incompatible shapes [3,4] and [3,5]

fn process(input: Tensor<Float32, [?, 784]>): Tensor<Float32, [?, 10]> {
    val w = randn([784, 10]);
    input @ w    // batch dim (?) propagated, inner dim (784) checked statically
}

val tensor = randn([1024, 1024]);
val other = tensor;          // tensor is MOVED into other
// println("{}", tensor);    // ERROR[E4001]: use of moved value `tensor`
println("{}", other);        // OK

fn consume(t: Tensor<Float32, [3, 3]>) {
    println("{}", t);
}

val data = randn([3, 3]);
consume(data);
// consume(data);   // ERROR[E4001]: use of moved value `data`

fn print_shape(t: &amp;Tensor<Float32, [?, 784]>) {
    println("Shape: {}", t.shape);
}

val input = randn([32, 784]);
print_shape(&amp;input);       // borrow, don't move
print_shape(&amp;input);       // still valid — input wasn't moved

fn scale(t: &amp;mut Tensor<Float32, [3, 3]>, factor: Float32) {
    // modify tensor in place
}

var weights = randn([3, 3]);
scale(&amp;mut weights, 2.0);
println("{}", weights);    // OK — mutable borrow has ended

var data = randn([10]);
val r1 = &data;
val r2 = &mut data;       // ERROR[E4003]: cannot borrow `data` as mutable
                           // because it is also borrowed as immutable
println("{}", r1);

fn first_element(v: &amp;Vec<Int32>): &amp;Int32 {
    &amp;v[0]   // lifetime of return value tied to lifetime of `v`
}

fn longest<'a>(a: &amp;'a String, b: &amp;'a String): &amp;'a String {
    if a.len() > b.len() { a } else { b }
}

fn dangling(): &String {
    val s = "hello".to_string();
    &s   // ERROR[E4005]: `s` does not live long enough
}        // `s` is dropped here

val a: Int32 = 42;
val b = a;          // copy — both a and b are valid
println("{} {}", a, b);

val s = "hello".to_string();
val t = s;          // move — only t is valid
// println("{}", s);   // ERROR

model Color: Copy, Clone {
    r: UInt8,
    g: UInt8,
    b: UInt8,
}

val red = Color { r: 255, g: 0, b: 0 };
val also_red = red;   // copy, not move
println("{}", red.r);  // OK

var t = randn([256, 256]);
val cpu_ref = &t;
val gpu_t = t.to_gpu();      // ERROR[E4007]: cannot move `t` to GPU while
                              // borrowed on CPU

val cpu_data = randn([1024]);
val gpu_data = cpu_data.to_gpu();    // cpu_data is moved
// println("{}", cpu_data);          // ERROR: use of moved value

val result = gpu_data.to_cpu();      // gpu_data is moved back
println("{}", result);

var data = randn([256, 256]);

// Phase 1: work on CPU
val norm = data.mean();
println("Mean: {}", norm);

// Phase 2: transfer to GPU (no outstanding borrows)
val gpu_data = data.to_gpu();
val result = gpu_data @ gpu_data;

model Trainer {
    model: NeuralNet,
    optimizer: Adam,
}

extend Trainer {
    fn train_epoch(&mut self, data: &DataLoader): Float32 {
        var total_loss = 0.0;

        for batch in data {
            val (inputs, targets) = batch;

            // model borrowed mutably through self
            val predictions = self.model.forward(inputs);
            val loss = cross_entropy(predictions, targets);

            total_loss += loss.item();

            loss.backward();
            self.optimizer.step();
            self.optimizer.zero_grad();
        }

        total_loss / data.len() as Float32
    }
}

enum Option<T> {
    Some(T),
    None,
}

fn find_index(items: &amp;Vec<String>, target: &amp;String): Option<Int32> {
    for i in 0..items.len() {
        if items[i] == target {
            return Some(i);
        }
    }
    None
}

fn main() {
    val names = vec!["Alice", "Bob", "Charlie"];

    match find_index(&amp;names, "Bob") {
        Some(idx) => println("Found at index {}", idx),
        None => println("Not found"),
    }
}

val opt: Option<Int32> = Some(42);

// Unwrap (panics if None)
val x = opt.unwrap();              // 42

// Unwrap with default
val y = opt.unwrap_or(0);          // 42
val z = None.unwrap_or(0);         // 0

// Map: transform the inner value
val doubled = opt.map(|x| x * 2); // Some(84)

// is_some / is_none
if opt.is_some() {
    println("Has a value");
}

// and_then: chain optional operations
val result = opt
    .map(|x| x + 1)
    .and_then(|x| if x > 0 { Some(x) } else { None });

enum Result<T, E> {
    Ok(T),
    Err(E),
}

fn parse_int(s: &amp;String): Result<Int32, String> {
    // parsing logic...
    if valid {
        Ok(parsed_value)
    } else {
        Err("invalid integer: " + s)
    }
}

fn read_config(path: String): Result<Config, IOError> {
    val file = File.open(path)?;       // propagate error
    val contents = file.read_all()?;    // propagate error
    val config = parse_toml(contents)?;
    Ok(config)
}

val ok: Result<Int32, String> = Ok(42);
val err: Result<Int32, String> = Err("oops");

// Unwrap (panics on Err)
val x = ok.unwrap();           // 42

// Unwrap with default
val y = err.unwrap_or(0);      // 0

// Map the success value
val doubled = ok.map(|x| x * 2);   // Ok(84)

// Map the error
val mapped_err = err.map_err(|e| IOError.new(e));

// is_ok / is_err
if ok.is_ok() {
    println("Success!");
}

// and_then: chain fallible operations
val result = ok
    .and_then(|x| if x > 0 { Ok(x) } else { Err("negative") });

fn process_file(path: String) {
    match File.open(path) {
        Ok(file) => {
            match file.read_all() {
                Ok(data) => println("Read {} bytes", data.len()),
                Err(e) => eprintln("Read error: {}", e),
            }
        }
        Err(e) => eprintln("Open error: {}", e),
    }
}

match load_model("model.axon") {
    Ok(model) => {
        println("Model loaded: {} parameters", model.param_count());
    }
    Err(IOError.NotFound(path)) => {
        eprintln("File not found: {}", path);
    }
    Err(IOError.PermissionDenied(path)) => {
        eprintln("Permission denied: {}", path);
    }
    Err(e) => {
        eprintln("Unexpected error: {}", e);
    }
}

// Without ?
fn load_data(path: String): Result<Vec<Float32>, IOError> {
    val file = match File.open(path) {
        Ok(f) => f,
        Err(e) => return Err(e),
    };
    val contents = match file.read_all() {
        Ok(c) => c,
        Err(e) => return Err(e),
    };
    parse_csv(contents)
}

// With ? — equivalent but cleaner
fn load_data(path: String): Result<Vec<Float32>, IOError> {
    val file = File.open(path)?;
    val contents = file.read_all()?;
    parse_csv(contents)
}

fn pipeline(path: String): Result<Model, Error> {
    val config = load_config(path)?;
    val data = load_dataset(&amp;config.data_path)?;
    val model = build_model(&amp;config)?;
    val trained = train(model, data)?;
    Ok(trained)
}

fn connect(host: String): Result<Connection, NetworkError> {
    // network errors are expected — caller decides what to do
}

fn get_element(v: &amp;Vec<Int32>, idx: Int32): Int32 {
    if idx < 0 || idx >= v.len() {
        panic("index out of bounds: {} (len: {})", idx, v.len());
    }
    v[idx]
}

enum ModelError {
    LoadFailed(String),
    ShapeMismatch { expected: Vec<Int32>, actual: Vec<Int32> },
    TrainingDiverged,
}

extend Display for ModelError {
    fn to_string(&amp;self): String {
        match self {
            ModelError.LoadFailed(path) => format("failed to load: {}", path),
            ModelError.ShapeMismatch { expected, actual } =>
                format("shape mismatch: expected {:?}, got {:?}", expected, actual),
            ModelError.TrainingDiverged => "training diverged (loss = NaN)".to_string(),
        }
    }
}

fn load_and_train(path: String): Result<Model, ModelError> {
    val model = load_model(path).map_err(|e| ModelError.LoadFailed(e.to_string()))?;
    train(model)
}

fn train_model(config: &amp;TrainConfig): Result<Model, ModelError> {
    val data = DataLoader.from_csv(&amp;config.data_path)
        .map_err(|e| ModelError.LoadFailed(e.to_string()))?;

    var model = NeuralNet.new(config.hidden_size);
    var optimizer = Adam.new(model.parameters(), lr: config.learning_rate);

    for epoch in 0..config.epochs {
        var epoch_loss = 0.0;

        for batch in &amp;data {
            val (inputs, targets) = batch;
            val predictions = model.forward(inputs);
            val loss = cross_entropy(predictions, targets);

            // Check for divergence
            if loss.item().is_nan() {
                return Err(ModelError.TrainingDiverged);
            }

            epoch_loss += loss.item();
            loss.backward();
            optimizer.step();
            optimizer.zero_grad();
        }

        println("Epoch {}: loss = {:.4}", epoch, epoch_loss / data.len() as Float32);
    }

    Ok(model)
}

mod math {
    pub fn square(x: Float64): Float64 {
        x * x
    }

    pub fn cube(x: Float64): Float64 {
        x * x * x
    }

    fn helper() {
        // private — not visible outside `math`
    }
}

fn main() {
    println("{}", math.square(5.0));   // 25.0
    println("{}", math.cube(3.0));     // 27.0
    // math.helper();                  // ERROR: `helper` is private
}

my_project/
├── Axon.toml
└── src/
    ├── main.axon         # crate root
    ├── model.axon         # mod model
    ├── data/
    │   ├── mod.axon       # mod data (directory module)
    │   ├── loader.axon    # mod data.loader
    │   └── transform.axon # mod data.transform
    └── utils.axon         # mod utils

mod model;
mod data;
mod utils;

fn main() {
    val net = model.build_network();
    val loader = data.loader.DataLoader.new("train.csv");
    utils.log("Training started");
}

pub mod loader;
pub mod transform;

mod network {
    pub model Layer {
        pub size: Int32,        // public field
        weights: Vec<Float32>,  // private field
    }

    pub fn new_layer(size: Int32): Layer {
        Layer {
            size,
            weights: Vec.new(),
        }
    }

    extend Layer {
        pub fn forward(&amp;self, input: &amp;Vec<Float32>): Vec<Float32> {
            // public method
        }

        fn init_weights(&amp;mut self) {
            // private method — internal use only
        }
    }
}

use std.collections.HashMap;
use std.io.{File, Read, Write};

fn main() {
    var map = HashMap.new();
    map.insert("key", 42);
}

// Absolute path
use std.math.sin;

// Nested imports
use std.collections.{Vec, HashMap, HashSet};

// Wildcard import (use sparingly)
use std.prelude.*;

// Aliased import
use std.collections.HashMap as Map;

mod internal {
    pub fn core_function(): Int32 { 42 }
}

// Re-export so users see `my_lib.core_function`
pub use internal.core_function;

[package]
name = "my_ml_project"
version = "0.2.1"
edition = "2026"
authors = ["Jane Doe <jane@example.com>"]
description = "A neural network toolkit"
license = "MIT"
repository = "https://github.com/jane/my_ml_project"

[dependencies]
axon-vision = "0.3.0"
axon-nlp = { version = "0.1.0", features = ["transformers"] }
axon-data = { git = "https://github.com/axon-lang/axon-data.git", branch = "main" }

[dev-dependencies]
axon-test = "0.1.0"

[build]
opt-level = 2
gpu = "cuda"

# From registry
axonc pkg add axon-vision
axonc pkg add axon-nlp --version 0.2.0

# Remove a dependency
axonc pkg remove axon-vision

use axon_vision.transforms.{resize, normalize};
use axon_nlp.tokenizer.BPETokenizer;

fn preprocess(image: Tensor<Float32, [?, ?, 3]>): Tensor<Float32, [?, 224, 224, 3]> {
    val resized = resize(image, [224, 224]);
    normalize(resized, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225])
}

my_ml_project/
├── Axon.toml
├── Axon.lock
├── src/
│   ├── main.axon           # entry point
│   ├── lib.axon             # library root (if building a library)
│   ├── model/
│   │   ├── mod.axon
│   │   ├── encoder.axon
│   │   └── decoder.axon
│   ├── data/
│   │   ├── mod.axon
│   │   └── preprocessing.axon
│   └── utils.axon
├── tests/
│   ├── test_model.axon
│   └── test_data.axon
├── benches/
│   └── bench_model.axon
└── examples/
    └── inference.axon

@gpu
fn vector_add(a: Tensor<Float32, [1024]>, b: Tensor<Float32, [1024]>): Tensor<Float32, [1024]> {
    a + b
}

@cpu
fn save_results(data: Tensor<Float32, [?, 10]>, path: String) {
    val file = File.create(path);
    file.write(data);
}

@device("cuda:0")
fn forward_gpu0(x: Tensor<Float32, [?, 784]>): Tensor<Float32, [?, 10]> {
    // executes on CUDA device 0
    val w = randn([784, 10]);
    x @ w
}

fn gpu_example() {
    // Create on CPU
    val cpu_tensor = randn([1024, 1024]);

    // Transfer to GPU
    val gpu_tensor = cpu_tensor.to_gpu();

    // Compute on GPU — fast!
    val result = gpu_tensor @ gpu_tensor;

    // Transfer back to CPU for I/O
    val cpu_result = result.to_cpu();
    println("{}", cpu_result);
}

val data = randn([256, 256]);
val gpu_data = data.to_gpu();   // data is moved
// println("{}", data);          // ERROR[E4001]: use of moved value `data`

val data = randn([256, 256]);
val backup = data.clone();
val gpu_data = data.to_gpu();
println("{}", backup);           // OK — backup is a separate copy

val cpu_a = randn([100]);
val gpu_b = randn([100]).to_gpu();
// val c = cpu_a + gpu_b;  // ERROR: device mismatch — cpu and gpu tensors

@gpu
fn init_weights(): Tensor<Float32, [784, 256]> {
    randn([784, 256])   // created directly on GPU — no transfer needed
}

# Compile for NVIDIA GPUs
axonc build model.axon --gpu cuda -O 3

# Compile for AMD GPUs
axonc build model.axon --gpu rocm -O 3

# Compile for Vulkan (cross-platform)
axonc build model.axon --gpu vulkan -O 3

use std.device.{Device, cuda};

fn main() {
    val dev0 = cuda(0);   // first GPU
    val dev1 = cuda(1);   // second GPU

    val a = randn([1024, 1024]).to_device(dev0);
    val b = randn([1024, 1024]).to_device(dev1);
}

fn train_multi_gpu(model: &mut NeuralNet, data: &DataLoader) {
    val devices = [cuda(0), cuda(1)];

    for batch in data {
        val (inputs, targets) = batch;

        // Split batch across devices
        val chunks = inputs.chunk(devices.len(), dim: 0);

        var losses = Vec.new();
        for i in 0..devices.len() {
            val chunk = chunks[i].to_device(devices[i]);
            val pred = model.forward(chunk);
            val loss = cross_entropy(pred, targets);
            losses.push(loss);
        }

        // Aggregate gradients
        val total_loss = losses.sum();
        total_loss.backward();
    }
}

use std.device;

fn main() {
    val count = device.gpu_count();
    println("Available GPUs: {}", count);

    for i in 0..count {
        val dev = device.cuda(i);
        println("  GPU {}: {} ({}MB)", i, dev.name(), dev.memory_mb());
    }
}

use std.device.cuda;

@gpu
fn matmul_gpu(
    a: Tensor<Float32, [?, ?]>,
    b: Tensor<Float32, [?, ?]>,
): Tensor<Float32, [?, ?]> {
    a @ b
}

fn main() {
    val size = 2048;

    // Create tensors on CPU
    val a = randn([size, size]);
    val b = randn([size, size]);

    // Transfer to GPU
    val ga = a.to_gpu();
    val gb = b.to_gpu();

    // GPU matrix multiply
    val gc = matmul_gpu(ga, gb);

    // Get result
    val c = gc.to_cpu();
    println("Result shape: {}", c.shape);
    println("Result[0][0]: {}", c[0][0]);
}

axonc build matmul.axon --gpu cuda -O 3 -o matmul
./matmul
# Result shape: [2048, 2048]
# Result[0][0]: 12.3456

axonc pkg new hello_tensor
cd hello_tensor

fn main() {
    // Create a 1D tensor from values
    val numbers = Tensor.from_vec([1.0, 2.0, 3.0, 4.0, 5.0], [5]);
    println("Numbers: {}", numbers);
    println("Shape:   {}", numbers.shape);
    println("Sum:     {}", numbers.sum());
    println("Mean:    {}", numbers.mean());
}

axonc pkg run
# Numbers: [1.0, 2.0, 3.0, 4.0, 5.0]
# Shape:   [5]
# Sum:     15.0
# Mean:    3.0

fn main() {
    // Zeros and ones
    val z: Tensor<Float32, [2, 3]> = zeros([2, 3]);
    println("Zeros:\n{}", z);

    val o: Tensor<Float32, [3]> = ones([3]);
    println("Ones: {}", o);

    // Random tensors
    val r = randn([2, 2]);     // normal distribution
    println("Random:\n{}", r);

    // Range
    val seq = arange(0, 5);
    println("Range: {}", seq);   // [0, 1, 2, 3, 4]

    // Identity matrix
    val eye = Tensor.eye(3);
    println("Identity:\n{}", eye);
}

fn main() {
    val a = Tensor.from_vec([1.0, 2.0, 3.0], [3]);
    val b = Tensor.from_vec([4.0, 5.0, 6.0], [3]);

    println("a + b = {}", a + b);    // [5.0, 7.0, 9.0]
    println("a * b = {}", a * b);    // [4.0, 10.0, 18.0]
    println("a * 2 = {}", a * 2.0);  // [2.0, 4.0, 6.0]

    // Math functions
    val x = Tensor.from_vec([0.0, 1.5708, 3.1416], [3]);
    println("sin(x) = {}", x.sin());
    println("exp(x) = {}", x.exp());
}

fn main() {
    val A: Tensor<Float32, [2, 3]> = Tensor.from_vec(
        [1.0, 2.0, 3.0,
         4.0, 5.0, 6.0], [2, 3]
    );

    val B: Tensor<Float32, [3, 2]> = Tensor.from_vec(
        [7.0,  8.0,
         9.0,  10.0,
         11.0, 12.0], [3, 2]
    );

    val C = A @ B;   // [2, 3] @ [3, 2] → [2, 2]
    println("A @ B =\n{}", C);
    println("Shape: {}", C.shape);
    // A @ B =
    // [[58.0, 64.0],
    //  [139.0, 154.0]]
}

fn main() {
    val t: Tensor<Float32, [2, 6]> = arange(0, 12).reshape([2, 6]);
    println("Original [2, 6]:\n{}", t);

    // Reshape
    val r = t.reshape([3, 4]);
    println("Reshaped [3, 4]:\n{}", r);

    val flat = t.reshape([12]);
    println("Flat: {}", flat);

    // Transpose
    val m: Tensor<Float32, [2, 3]> = Tensor.from_vec(
        [1.0, 2.0, 3.0,
         4.0, 5.0, 6.0], [2, 3]
    );
    val mt = m.transpose();
    println("Transposed [3, 2]:\n{}", mt);
}

fn main() {
    val data: Tensor<Float32, [3, 4]> = Tensor.from_vec(
        [1.0, 2.0, 3.0, 4.0,
         5.0, 6.0, 7.0, 8.0,
         9.0, 10.0, 11.0, 12.0], [3, 4]
    );

    println("Sum (all):   {}", data.sum());          // 78.0
    println("Mean (all):  {}", data.mean());         // 6.5
    println("Max (all):   {}", data.max());          // 12.0
    println("Sum (dim 0): {}", data.sum(dim: 0));    // [15.0, 18.0, 21.0, 24.0]
    println("Sum (dim 1): {}", data.sum(dim: 1));    // [10.0, 26.0, 42.0]
}

fn normalize(data: Tensor<Float32, [?, ?]>): Tensor<Float32, [?, ?]> {
    val mean = data.mean(dim: 0);
    val std = data.std(dim: 0);
    (data - mean) / std
}

fn main() {
    // Simulate a dataset: 100 samples, 4 features
    val dataset = randn([100, 4]);

    println("Before normalization:");
    println("  Mean: {}", dataset.mean(dim: 0));

    val normed = normalize(dataset);

    println("After normalization:");
    println("  Mean: {}", normed.mean(dim: 0));   // ≈ [0, 0, 0, 0]
    println("  Std:  {}", normed.std(dim: 0));    // ≈ [1, 1, 1, 1]
}

fn generate_data(n: Int32): (Tensor<Float32, [?, 1]>, Tensor<Float32, [?, 1]>) {
    // True parameters: y = 3.5x + 2.0 + noise
    val x = randn([n, 1]);
    val noise = randn([n, 1]) * 0.3;
    val y = x * 3.5 + 2.0 + noise;
    (x, y)
}

model LinearRegression {
    weight: Tensor<Float32, [1, 1]>,
    bias: Tensor<Float32, [1]>,
}

extend LinearRegression {
    fn new(): LinearRegression {
        LinearRegression {
            weight: randn([1, 1]),
            bias: zeros([1]),
        }
    }

    fn forward(&amp;self, x: Tensor<Float32, [?, 1]>): Tensor<Float32, [?, 1]> {
        x @ self.weight + self.bias
    }
}

fn mse_loss(
    predictions: Tensor<Float32, [?, 1]>,
    targets: Tensor<Float32, [?, 1]>,
): Tensor<Float32, []> {
    val diff = predictions - targets;
    (diff * diff).mean()
}

use std.autograd.GradTensor;
use std.optim.SGD;

fn train() {
    // Generate training data
    val (x_train, y_train) = generate_data(200);

    // Initialize model
    var net = LinearRegression.new();

    // Optimizer: SGD with learning rate 0.01
    var optimizer = SGD.new(
        [&net.weight, &net.bias],
        lr: 0.01,
    );

    // Training loop
    for epoch in 0..100 {
        // Forward pass
        val predictions = net.forward(x_train);

        // Compute loss
        val loss = mse_loss(predictions, y_train);

        // Backward pass — compute gradients
        loss.backward();

        // Update parameters
        optimizer.step();
        optimizer.zero_grad();

        if epoch % 10 == 0 {
            println("Epoch {}: loss = {:.4}", epoch, loss.item());
        }
    }

    // Print learned parameters
    println("Learned weight: {:.4} (true: 3.5)", net.weight.item());
    println("Learned bias:   {:.4} (true: 2.0)", net.bias.item());
}

fn evaluate(net: &LinearRegression) {
    // Generate test data
    val (x_test, y_test) = generate_data(50);

    // Predict
    val predictions = net.forward(x_test);

    // Compute test loss
    val test_loss = mse_loss(predictions, y_test);
    println("Test MSE: {:.4}", test_loss.item());

    // Print a few predictions
    println("\nSample predictions:");
    println("  x       | predicted | actual");
    println("  --------|-----------|-------");
    for i in 0..5 {
        println("  {:.4}  | {:.4}    | {:.4}",
            x_test[i].item(),
            predictions[i].item(),
            y_test[i].item()
        );
    }
}

use std.autograd.GradTensor;
use std.optim.SGD;

model LinearRegression {
    weight: Tensor<Float32, [1, 1]>,
    bias: Tensor<Float32, [1]>,
}

extend LinearRegression {
    fn new(): LinearRegression {
        LinearRegression {
            weight: randn([1, 1]),
            bias: zeros([1]),
        }
    }

    fn forward(&amp;self, x: Tensor<Float32, [?, 1]>): Tensor<Float32, [?, 1]> {
        x @ self.weight + self.bias
    }
}

fn mse_loss(
    predictions: Tensor<Float32, [?, 1]>,
    targets: Tensor<Float32, [?, 1]>,
): Tensor<Float32, []> {
    val diff = predictions - targets;
    (diff * diff).mean()
}

fn main() {
    println("=== Linear Regression in Axon ===\n");

    // Data
    val (x_train, y_train) = generate_data(200);

    // Model
    var net = LinearRegression.new();
    var optimizer = SGD.new(
        [&amp;net.weight, &amp;net.bias],
        lr: 0.01,
    );

    // Train
    for epoch in 0..200 {
        val pred = net.forward(x_train);
        val loss = mse_loss(pred, y_train);
        loss.backward();
        optimizer.step();
        optimizer.zero_grad();

        if epoch % 50 == 0 {
            println("Epoch {:>3}: loss = {:.6}", epoch, loss.item());
        }
    }

    println("\nFinal parameters:");
    println("  weight = {:.4} (true: 3.5)", net.weight.item());
    println("  bias   = {:.4} (true: 2.0)", net.bias.item());

    // Evaluate
    val (x_test, y_test) = generate_data(50);
    val test_pred = net.forward(x_test);
    val test_loss = mse_loss(test_pred, y_test);
    println("\nTest MSE: {:.6}", test_loss.item());
}

fn generate_data(n: Int32): (Tensor<Float32, [?, 1]>, Tensor<Float32, [?, 1]>) {
    val x = randn([n, 1]);
    val noise = randn([n, 1]) * 0.3;
    val y = x * 3.5 + 2.0 + noise;
    (x, y)
}

axonc build linear_reg.axon -O 2 -o linear_reg
./linear_reg
# === Linear Regression in Axon ===
#
# Epoch   0: loss = 14.832901
# Epoch  50: loss = 0.129384
# Epoch 100: loss = 0.091203
# Epoch 150: loss = 0.089847
#
# Final parameters:
#   weight = 3.4821 (true: 3.5)
#   bias   = 1.9934 (true: 2.0)
#
# Test MSE: 0.092145

use std.data.DataLoader;
use std.transforms.{normalize, to_tensor};

fn load_mnist(): (DataLoader, DataLoader) {
    val train_loader = DataLoader.from_csv("data/mnist_train.csv")
        .batch_size(64)
        .shuffle(true)
        .transform(|img| {
            val tensor = to_tensor(img, [1, 28, 28]);
            normalize(tensor, mean: [0.1307], std: [0.3081])
        });

    val test_loader = DataLoader.from_csv("data/mnist_test.csv")
        .batch_size(256)
        .shuffle(false)
        .transform(|img| {
            val tensor = to_tensor(img, [1, 28, 28]);
            normalize(tensor, mean: [0.1307], std: [0.3081])
        });

    (train_loader, test_loader)
}

use std.nn.{Conv2d, Linear, MaxPool2d, Module, Sequential};

model MNISTNet {
    conv1: Conv2d,
    conv2: Conv2d,
    pool: MaxPool2d,
    fc1: Linear<9216, 128>,
    fc2: Linear<128, 10>,
}

extend MNISTNet {
    fn new(): MNISTNet {
        MNISTNet {
            conv1: Conv2d.new(in_channels: 1, out_channels: 32, kernel_size: 3, padding: 1),
            conv2: Conv2d.new(in_channels: 32, out_channels: 64, kernel_size: 3, padding: 1),
            pool: MaxPool2d.new(kernel_size: 2, stride: 2),
            fc1: Linear.new(),
            fc2: Linear.new(),
        }
    }
}

extend Module for MNISTNet {
    fn forward(&amp;self, x: Tensor<Float32, [?, 1, 28, 28]>): Tensor<Float32, [?, 10]> {
        // Conv block 1: [?, 1, 28, 28] → [?, 32, 14, 14]
        val h = self.conv1.forward(x);
        val h = relu(h);
        val h = self.pool.forward(h);

        // Conv block 2: [?, 32, 14, 14] → [?, 64, 7, 7]
        val h = self.conv2.forward(h);
        val h = relu(h);
        val h = self.pool.forward(h);

        // Flatten: [?, 64, 7, 7] → [?, 3136]
        val batch_size = h.shape[0];
        val h = h.reshape([batch_size, 3136]);

        // Fully connected layers
        val h = relu(self.fc1.forward(h));
        self.fc2.forward(h)
    }
}

use std.optim.Adam;
use std.loss.cross_entropy;
use std.metrics.accuracy;

fn train_epoch(
    net: &mut MNISTNet,
    data: &DataLoader,
    optimizer: &mut Adam,
): (Float32, Float32) {
    var total_loss = 0.0;
    var correct = 0;
    var total = 0;

    for batch in data {
        val (images, labels) = batch;

        // Forward
        val logits = net.forward(images);
        val loss = cross_entropy(logits, labels);

        // Track metrics
        total_loss += loss.item();
        val predicted = logits.argmax(dim: 1);
        correct += (predicted == labels).sum().item() as Int32;
        total += labels.shape[0];

        // Backward
        loss.backward();
        optimizer.step();
        optimizer.zero_grad();
    }

    val avg_loss = total_loss / data.num_batches() as Float32;
    val acc = correct as Float32 / total as Float32;
    (avg_loss, acc)
}

fn evaluate(net: &MNISTNet, data: &DataLoader): (Float32, Float32) {
    var total_loss = 0.0;
    var correct = 0;
    var total = 0;

    for batch in data {
        val (images, labels) = batch;
        val logits = net.forward(images);
        val loss = cross_entropy(logits, labels);

        total_loss += loss.item();
        val predicted = logits.argmax(dim: 1);
        correct += (predicted == labels).sum().item() as Int32;
        total += labels.shape[0];
    }

    val avg_loss = total_loss / data.num_batches() as Float32;
    val acc = correct as Float32 / total as Float32;
    (avg_loss, acc)
}

use std.nn.{Conv2d, Linear, MaxPool2d, Module};
use std.optim.Adam;
use std.loss.cross_entropy;
use std.data.DataLoader;
use std.transforms.{normalize, to_tensor};

fn main() {
    println("=== MNIST Classifier ===\n");

    // Load data
    val (train_loader, test_loader) = load_mnist();
    println("Train: {} samples", train_loader.len());
    println("Test:  {} samples\n", test_loader.len());

    // Create model and optimizer
    var net = MNISTNet.new();
    var optimizer = Adam.new(
        net.parameters(),
        lr: 0.001,
    );

    // Training
    val epochs = 10;
    for epoch in 0..epochs {
        val (train_loss, train_acc) = train_epoch(&mut net, &train_loader, &mut optimizer);
        val (test_loss, test_acc) = evaluate(&net, &test_loader);

        println("Epoch {:>2}/{} | Train Loss: {:.4} Acc: {:.2}% | Test Loss: {:.4} Acc: {:.2}%",
            epoch + 1, epochs,
            train_loss, train_acc * 100.0,
            test_loss, test_acc * 100.0,
        );
    }

    // Final evaluation
    val (_, final_acc) = evaluate(&net, &test_loader);
    println("\nFinal test accuracy: {:.2}%", final_acc * 100.0);
}

=== MNIST Classifier ===

Train: 60000 samples
Test:  10000 samples

Epoch  1/10 | Train Loss: 0.2134 Acc: 93.41% | Test Loss: 0.0712 Acc: 97.82%
Epoch  2/10 | Train Loss: 0.0583 Acc: 98.19% | Test Loss: 0.0498 Acc: 98.41%
...
Epoch 10/10 | Train Loss: 0.0089 Acc: 99.71% | Test Loss: 0.0312 Acc: 99.12%

Final test accuracy: 99.12%

fn main() {
    var net = MNISTNet.new().to_gpu();
    var optimizer = Adam.new(net.parameters(), lr: 0.001);

    for epoch in 0..10 {
        for batch in &train_loader {
            val (images, labels) = batch;
            val images = images.to_gpu();
            val labels = labels.to_gpu();

            val logits = net.forward(images);
            val loss = cross_entropy(logits, labels);
            loss.backward();
            optimizer.step();
            optimizer.zero_grad();
        }
    }
}

axonc build mnist.axon --gpu cuda -O 3 -o mnist

use std.export.save;

// After training
save(&net, "mnist_model.axon");
println("Model saved!");

// Load later
use std.export.load;
val loaded_net: MNISTNet = load("mnist_model.axon");

fn scaled_dot_product_attention(
    query: Tensor<Float32, [?, ?, ?]>,    // [batch, seq_len, d_k]
    key: Tensor<Float32, [?, ?, ?]>,      // [batch, seq_len, d_k]
    value: Tensor<Float32, [?, ?, ?]>,    // [batch, seq_len, d_v]
): Tensor<Float32, [?, ?, ?]> {
    val d_k = query.shape[2] as Float32;

    // Attention scores: Q @ K^T / sqrt(d_k)
    val scores = (query @ key.transpose()) / d_k.sqrt();

    // Softmax over the last dimension
    val weights = softmax(scores, dim: 2);

    // Weighted sum of values
    weights @ value
}

use std.nn.Linear;

model MultiHeadAttention {
    num_heads: Int32,
    d_model: Int32,
    d_k: Int32,
    w_query: Linear<512, 512>,
    w_key: Linear<512, 512>,
    w_value: Linear<512, 512>,
    w_output: Linear<512, 512>,
}

extend MultiHeadAttention {
    fn new(d_model: Int32, num_heads: Int32): MultiHeadAttention {
        val d_k = d_model / num_heads;
        MultiHeadAttention {
            num_heads,
            d_model,
            d_k,
            w_query: Linear.new(),
            w_key: Linear.new(),
            w_value: Linear.new(),
            w_output: Linear.new(),
        }
    }

    fn forward(
        &amp;self,
        query: Tensor<Float32, [?, ?, 512]>,
        key: Tensor<Float32, [?, ?, 512]>,
        value: Tensor<Float32, [?, ?, 512]>,
    ): Tensor<Float32, [?, ?, 512]> {
        val batch_size = query.shape[0];
        val seq_len = query.shape[1];

        // Project Q, K, V
        val q = self.w_query.forward(query);
        val k = self.w_key.forward(key);
        val v = self.w_value.forward(value);

        // Reshape to [batch, num_heads, seq_len, d_k]
        val q = q.reshape([batch_size, seq_len, self.num_heads, self.d_k])
                  .permute([0, 2, 1, 3]);
        val k = k.reshape([batch_size, seq_len, self.num_heads, self.d_k])
                  .permute([0, 2, 1, 3]);
        val v = v.reshape([batch_size, seq_len, self.num_heads, self.d_k])
                  .permute([0, 2, 1, 3]);

        // Attention per head
        val attn = scaled_dot_product_attention(q, k, v);

        // Concatenate heads
        val concat = attn.permute([0, 2, 1, 3])
                         .reshape([batch_size, seq_len, self.d_model]);

        // Final projection
        self.w_output.forward(concat)
    }
}

model FeedForward {
    linear1: Linear<512, 2048>,
    linear2: Linear<2048, 512>,
}

extend FeedForward {
    fn new(d_model: Int32, d_ff: Int32): FeedForward {
        FeedForward {
            linear1: Linear.new(),
            linear2: Linear.new(),
        }
    }

    fn forward(&amp;self, x: Tensor<Float32, [?, ?, 512]>): Tensor<Float32, [?, ?, 512]> {
        val h = gelu(self.linear1.forward(x));
        self.linear2.forward(h)
    }
}

use std.nn.LayerNorm;

model TransformerBlock {
    attention: MultiHeadAttention,
    feed_forward: FeedForward,
    norm1: LayerNorm,
    norm2: LayerNorm,
}

extend TransformerBlock {
    fn new(d_model: Int32, num_heads: Int32, d_ff: Int32): TransformerBlock {
        TransformerBlock {
            attention: MultiHeadAttention.new(d_model, num_heads),
            feed_forward: FeedForward.new(d_model, d_ff),
            norm1: LayerNorm.new(d_model),
            norm2: LayerNorm.new(d_model),
        }
    }

    fn forward(&amp;self, x: Tensor<Float32, [?, ?, 512]>): Tensor<Float32, [?, ?, 512]> {
        // Self-attention + residual + norm
        val attn_out = self.attention.forward(x, x, x);
        val h = self.norm1.forward(x + attn_out);

        // Feed-forward + residual + norm
        val ff_out = self.feed_forward.forward(h);
        self.norm2.forward(h + ff_out)
    }
}

fn positional_encoding(seq_len: Int32, d_model: Int32): Tensor<Float32, [?, ?]> {
    val pe = zeros([seq_len, d_model]);

    for pos in 0..seq_len {
        for i in 0..(d_model / 2) {
            val angle = pos as Float32 / (10000.0).pow(2.0 * i as Float32 / d_model as Float32);
            pe[pos][2 * i] = angle.sin();
            pe[pos][2 * i + 1] = angle.cos();
        }
    }

    pe
}

use std.nn.{Embedding, Linear, Module};

model TransformerEncoder {
    embedding: Embedding,
    layers: Vec<TransformerBlock>,
    classifier: Linear<512, 10>,
    d_model: Int32,
}

extend TransformerEncoder {
    fn new(
        vocab_size: Int32,
        d_model: Int32,
        num_heads: Int32,
        num_layers: Int32,
        d_ff: Int32,
        num_classes: Int32,
    ): TransformerEncoder {
        var layers = Vec.new();
        for _ in 0..num_layers {
            layers.push(TransformerBlock.new(d_model, num_heads, d_ff));
        }

        TransformerEncoder {
            embedding: Embedding.new(vocab_size, d_model),
            layers,
            classifier: Linear.new(),
            d_model,
        }
    }
}

extend Module for TransformerEncoder {
    fn forward(&amp;self, tokens: Tensor<Int64, [?, ?]>): Tensor<Float32, [?, 10]> {
        val seq_len = tokens.shape[1];

        // Token embedding + positional encoding
        val x = self.embedding.forward(tokens);
        val pe = positional_encoding(seq_len, self.d_model);
        var h = x + pe;

        // Pass through transformer blocks
        for layer in &amp;self.layers {
            h = layer.forward(h);
        }

        // Classification: use [CLS] token (first position)
        val cls = h[.., 0, ..];   // [batch, d_model]
        self.classifier.forward(cls)
    }
}

use std.optim.AdamW;
use std.loss.cross_entropy;

fn main() {
    println("=== Transformer Encoder ===\n");

    // Hyperparameters
    val vocab_size = 10000;
    val d_model = 512;
    val num_heads = 8;
    val num_layers = 6;
    val d_ff = 2048;
    val num_classes = 10;

    // Create model
    var net = TransformerEncoder.new(
        vocab_size, d_model, num_heads, num_layers, d_ff, num_classes,
    );

    var optimizer = AdamW.new(
        net.parameters(),
        lr: 0.0001,
        weight_decay: 0.01,
    );

    println("Model: {} parameters", net.param_count());
    println("Config: d_model={}, heads={}, layers={}\n", d_model, num_heads, num_layers);

    // Training loop
    val epochs = 20;
    for epoch in 0..epochs {
        var total_loss = 0.0;
        var num_batches = 0;

        for batch in &train_loader {
            val (tokens, labels) = batch;

            val logits = net.forward(tokens);
            val loss = cross_entropy(logits, labels);

            loss.backward();
            optimizer.step();
            optimizer.zero_grad();

            total_loss += loss.item();
            num_batches += 1;
        }

        val avg_loss = total_loss / num_batches as Float32;
        println("Epoch {:>2}/{}: loss = {:.4}", epoch + 1, epochs, avg_loss);
    }
}

use std.nn.{TransformerEncoder as TE, MultiHeadAttention};

fn main() {
    // One-liner transformer encoder
    val encoder = TE.new(
        d_model: 512,
        num_heads: 8,
        num_layers: 6,
        d_ff: 2048,
        dropout: 0.1,
    );

    val input: Tensor<Float32, [?, 128, 512]> = randn([32, 128, 512]);
    val output = encoder.forward(input);
    println("Output shape: {}", output.shape);   // [32, 128, 512]
}

Input Tokens
     │
     ▼
┌─────────────┐
│  Embedding   │
│  + Pos Enc   │
└─────┬───────┘
      │
      ▼  ×N layers
┌─────────────────────┐
│  Multi-Head Attn    │
│  + Residual + Norm  │
├─────────────────────┤
│  Feed-Forward       │
│  + Residual + Norm  │
└─────────┬───────────┘
          │
          ▼
┌─────────────┐
│  Classifier  │
│  (Linear)    │
└─────────────┘
          │
          ▼
      Logits [?, num_classes]

model Point {
    x: Float64,
    y: Float64,
}

fn distance(a: Point, b: Point): Float64 {
    val dx = a.x - b.x;
    val dy = a.y - b.y;
    return (dx * dx + dy * dy).sqrt();
}

fn main() {
    val origin = Point { x: 0.0, y: 0.0 };
    val target = Point { x: 3.0, y: 4.0 };
    println(distance(origin, target));  // 5.0
}

model ModelConfig {
    learning_rate: Float64,
    batch_size: Int32,
    epochs: Int32,
}

extend ModelConfig {
    fn default(): ModelConfig {
        return ModelConfig {
            learning_rate: 0.001,
            batch_size: 32,
            epochs: 10,
        };
    }

    fn with_lr(self, lr: Float64): ModelConfig {
        return ModelConfig {
            learning_rate: lr,
            batch_size: self.batch_size,
            epochs: self.epochs,
        };
    }
}

model LinearLayer {
    weights: Tensor<Float32, [_, _]>,
    bias: Tensor<Float32, [_]>,
}

extend LinearLayer {
    fn forward(self, input: Tensor<Float32, [_, _]>): Tensor<Float32, [_, _]> {
        return input @ self.weights + self.bias;
    }
}

enum Activation {
    ReLU,
    Sigmoid,
    Tanh,
    LeakyReLU(Float64),  // variant with data
}

fn apply_activation(x: Float64, act: Activation): Float64 {
    match act {
        Activation.ReLU => if x > 0.0 { x } else { 0.0 },
        Activation.Sigmoid => 1.0 / (1.0 + (-x).exp()),
        Activation.Tanh => x.tanh(),
        Activation.LeakyReLU(alpha) => if x > 0.0 { x } else { alpha * x },
    }
}

enum Device {
    CPU,
    GPU(Int32),  // GPU with device index
}

fn device_name(d: Device): String {
    match d {
        Device.CPU => "cpu",
        Device.GPU(idx) => format("cuda:{}", idx),
    }
}

model DataBatch {
    images: Tensor<Float32, [_, 3, 224, 224]>,
    labels: Tensor<Int64, [_]>,
}

fn process_batch(batch: DataBatch) {
    // `batch` is moved here — caller can no longer use it
    val predictions = model.forward(batch.images);
    val loss = cross_entropy(predictions, batch.labels);
}

fn inspect_batch(batch: &DataBatch) {
    println(batch.images.shape());
    println(batch.labels.shape());
    // batch is borrowed — caller retains ownership
}

enum Option<T> {
    Some(T),
    None,
}

fn find_max(data: Tensor<Float32, [_]>): Option<Float32> {
    if data.len() == 0 {
        return None;
    }
    return Some(data.max());
}

fn main() {
    val values = tensor([1.0, 5.0, 3.0, 9.0, 2.0]);

    match find_max(values) {
        Some(max) => println("Max value: {}", max),
        None => println("Empty tensor"),
    }
}

val maybe_value: Option<Float64> = Some(42.0);

// unwrap_or: provide a default
val value = maybe_value.unwrap_or(0.0);

// map: transform the inner value
val doubled = maybe_value.map(|x| x * 2.0);

// is_some / is_none: check presence
if maybe_value.is_some() {
    println("Got a value!");
}

enum Result<T, E> {
    Ok(T),
    Err(E),
}

fn load_model(path: String): Result<Model, String> {
    if !file_exists(path) {
        return Err("Model file not found: " + path);
    }
    val data = read_file(path)?;  // ? propagates errors
    return Ok(parse_model(data));
}

fn main() {
    match load_model("weights.axon") {
        Ok(model) => println("Loaded model with {} params", model.param_count()),
        Err(e) => println("Error: {}", e),
    }
}

fn train_pipeline(config_path: String): Result<Float64, String> {
    val config = load_config(config_path)?;     // returns Err early if this fails
    val data = load_dataset(config.data_path)?;  // same here
    val model = build_model(config)?;             // and here

    val final_loss = train(model, data, config.epochs)?;
    return Ok(final_loss);
}

enum TrainingError {
    DataNotFound(String),
    InvalidShape { expected: Shape, actual: Shape },
    ConvergenceFailure { epoch: Int32, loss: Float64 },
    OutOfMemory,
}

fn train(model: Model, data: Dataset): Result<Model, TrainingError> {
    if data.shape() != model.expected_input_shape() {
        return Err(TrainingError.InvalidShape {
            expected: model.expected_input_shape(),
            actual: data.shape(),
        });
    }
    // ... training logic ...
    return Ok(model);
}

// Option → Result: provide an error message for the None case
val value: Result<Float64, String> = maybe_value.ok_or("value was missing");

// Result → Option: discard the error info
val maybe: Option<Float64> = result.ok();

fn assert_valid_shape(t: Tensor<Float32, [_, _]>) {
    if t.shape()[0] == 0 {
        panic("tensor must have at least one row");
    }
}

# PyTorch
import torch

x = torch.zeros(3, 4)
y = torch.ones(5)
z = torch.randn(128, 256)
a = torch.tensor([1.0, 2.0, 3.0])
e = torch.eye(4)
r = torch.arange(0, 10)

// Axon
val x = zeros([3, 4]);
val y = ones([5]);
val z = randn([128, 256]);
val a = Tensor.from_vec([1.0, 2.0, 3.0], [3]);
val e = Tensor.eye(4);
val r = arange(0, 10);

# PyTorch
c = a + b
c = a * b
c = a @ b          # matmul
c = torch.matmul(a, b)
m = x.mean(dim=0)
s = x.sum(dim=1)
r = x.reshape(3, 4)
t = x.T             # transpose

// Axon
val c = a + b;
val c = a * b;
val c = a @ b;          // matmul (same!)
// no functional form needed
val m = x.mean(dim: 0);
val s = x.sum(dim: 1);
val r = x.reshape([3, 4]);
val t = x.transpose();

# PyTorch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

model = MyModel()

// Axon
use std.nn.{Linear, Module};

model MyModel {
    fc1: Linear<784, 256>,
    fc2: Linear<256, 128>,
    fc3: Linear<128, 10>,
}

extend MyModel {
    fn new(): MyModel {
        MyModel {
            fc1: Linear.new(),
            fc2: Linear.new(),
            fc3: Linear.new(),
        }
    }
}

extend Module for MyModel {
    fn forward(&amp;self, x: Tensor<Float32, [?, 784]>): Tensor<Float32, [?, 10]> {
        val h = relu(self.fc1.forward(x));
        val h = relu(self.fc2.forward(h));
        self.fc3.forward(h)
    }
}

val model = MyModel.new();

# PyTorch
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    for inputs, targets in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}: loss = {loss.item():.4f}")

// Axon
use std.optim.Adam;
use std.loss.cross_entropy;

var optimizer = Adam.new(model.parameters(), lr: 0.001);

for epoch in 0..10 {
    for batch in &dataloader {
        val (inputs, targets) = batch;
        val outputs = model.forward(inputs);
        val loss = cross_entropy(outputs, targets);
        loss.backward();
        optimizer.step();
        optimizer.zero_grad();
    }
    println("Epoch {}: loss = {:.4}", epoch, loss.item());
}

# PyTorch
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2)
self.bn = nn.BatchNorm2d(32)
self.dropout = nn.Dropout(0.5)

// Axon
self.conv1 = Conv2d.new(in_channels: 1, out_channels: 32, kernel_size: 3, padding: 1);
self.pool = MaxPool2d.new(kernel_size: 2, stride: 2);
self.bn = BatchNorm.new(32);
self.dropout = Dropout.new(rate: 0.5);

# PyTorch
lstm = nn.LSTM(input_size=128, hidden_size=256, num_layers=2, batch_first=True)
attention = nn.MultiheadAttention(embed_dim=512, num_heads=8)
encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=512, nhead=8),
    num_layers=6
)

// Axon
val lstm = LSTM.new(input_size: 128, hidden_size: 256, num_layers: 2);
val attention = MultiHeadAttention.new(d_model: 512, num_heads: 8);
val encoder = TransformerEncoder.new(
    d_model: 512,
    num_heads: 8,
    num_layers: 6,
    d_ff: 2048,
    dropout: 0.1,
);

# PyTorch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
x = x.to(device)

# Multi-GPU
model = nn.DataParallel(model)

// Axon
use std.device.{cuda, cpu};

val model = MyModel.new().to_gpu();
val x = x.to_gpu();

// Or with explicit device
val dev = cuda(0);
val x = x.to_device(dev);

# PyTorch
x = torch.randn(3, requires_grad=True)
y = x * 2 + 1
y.sum().backward()
print(x.grad)

with torch.no_grad():
    prediction = model(x)

// Axon
use std.autograd.{GradTensor, no_grad};

val x = GradTensor.new(randn([3]));
val y = x * 2.0 + 1.0;
y.sum().backward();
println("{}", x.grad());

no_grad(|| {
    val prediction = model.forward(x);
});

# PyTorch
from torch.utils.data import DataLoader, TensorDataset

dataset = TensorDataset(x_train, y_train)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

// Axon
use std.data.DataLoader;

val loader = DataLoader.new(x_train, y_train)
    .batch_size(64)
    .shuffle(true);

# PyTorch
torch.save(model.state_dict(), "model.pt")
model.load_state_dict(torch.load("model.pt"))

# ONNX export
torch.onnx.export(model, dummy_input, "model.onnx")

// Axon
use std.export.{save, load, export_onnx};

save(&model, "model.axon");
val model: MyModel = load("model.axon");

// ONNX export
val dummy_input = randn([1, 784]);
export_onnx(&model, dummy_input, "model.onnx");

# Python
name = "Alice"
age = 30
scores = [95, 87, 92]

// Axon
val name = "Alice";
val age = 30;
val scores = vec![95, 87, 92];

# Python
def add(a: int, b: int) -> int:
    return a + b

def greet(name: str):
    print(f"Hello, {name}!")

// Axon
fn add(a: Int32, b: Int32): Int32 {
    a + b    // implicit return (last expression)
}

fn greet(name: String) {
    println("Hello, {}!", name);
}

# Python
if score >= 90:
    grade = "A"
elif score >= 70:
    grade = "B"
else:
    grade = "C"

// Axon — if is an expression!
val grade = if score >= 90 {
    "A"
} else if score >= 70 {
    "B"
} else {
    "C"
};

# Python
for i in range(10):
    print(i)

for item in items:
    process(item)

while condition:
    do_work()

// Axon
for i in 0..10 {
    println("{}", i);
}

for item in items {
    process(item);
}

while condition {
    do_work();
}

# Python 3.10+
match command:
    case "quit":
        exit()
    case "hello":
        print("Hi!")
    case _:
        print("Unknown")

// Axon
match command {
    "quit" => exit(),
    "hello" => println("Hi!"),
    _ => println("Unknown"),
}

# Python
class Point:
    def __init__(self, x: float, y: float):
        self.x = x
        self.y = y

    def distance(self, other: 'Point') -> float:
        return ((self.x - other.x)**2 + (self.y - other.y)**2)**0.5

    def __str__(self):
        return f"({self.x}, {self.y})"

// Axon
model Point {
    x: Float64,
    y: Float64,
}

extend Point {
    fn new(x: Float64, y: Float64): Point {
        Point { x, y }
    }

    fn distance(&self, other: &Point): Float64 {
        val dx = self.x - other.x;
        val dy = self.y - other.y;
        (dx * dx + dy * dy).sqrt()
    }
}

extend Display for Point {
    fn to_string(&self): String {
        format("({}, {})", self.x, self.y)
    }
}

# Python
try:
    f = open("config.toml")
    data = f.read()
    config = parse(data)
except FileNotFoundError:
    print("Config not found")
except ParseError as e:
    print(f"Parse error: {e}")

// Axon
match File.open("config.toml") {
    Ok(file) => {
        match file.read_all() {
            Ok(data) => {
                val config = parse(data);
                println("Loaded config");
            }
            Err(e) => eprintln("Read error: {}", e),
        }
    }
    Err(e) => eprintln("Config not found: {}", e),
}

// Or more concisely with ?
fn load_config(): Result<Config, IOError> {
    val file = File.open("config.toml")?;
    val data = file.read_all()?;
    parse(data)
}

# Python (NumPy)
import numpy as np

a = np.zeros((3, 4))
b = np.random.randn(3, 4)
c = a + b
d = np.dot(a, b.T)
mean = np.mean(c, axis=0)

// Axon
val a = zeros([3, 4]);
val b = randn([3, 4]);
val c = a + b;
val d = a @ b.transpose();
val mean = c.mean(dim: 0);

# Python
squares = [x**2 for x in range(10)]
evens = [x for x in numbers if x % 2 == 0]

// Axon (iterator methods)
val squares: Vec<Int32> = (0..10).map(|x| x * x).collect();
val evens: Vec<Int32> = numbers.iter().filter(|x| x % 2 == 0).collect();

# Python — file: math_utils.py
def square(x):
    return x * x

# main.py
from math_utils import square

// Axon — file: math_utils.axon
pub fn square(x: Float64): Float64 {
    x * x
}

// main.axon
mod math_utils;
use math_utils.square;

// Rust
fn add(x: i32, y: i32) -> i32 {
    x + y
}

// Axon — explicit return required, types use PascalCase
fn add(x: Int32, y: Int32): Int32 {
    return x + y;
}

// Rust
let x = 5;          // immutable
let mut y = 10;     // mutable

// Axon — val for immutable, var for mutable
val x = 5;
var y = 10;

fn take_ownership(data: Tensor<Float32, [_]>) {
    // `data` is owned here — moved from caller
    println(data.sum());
}
// `data` is dropped here

fn borrow_data(data: &amp;Tensor<Float32, [_]>) {
    // `data` is borrowed — caller keeps ownership
    println(data.sum());
}

// Axon — tensors are first-class citizens
val x: Tensor<Float32, [3, 3]> = tensor([[1.0, 2.0, 3.0],
                                          [4.0, 5.0, 6.0],
                                          [7.0, 8.0, 9.0]]);

// Shape-checked matrix multiply (compiler verifies dimensions)
val y: Tensor<Float32, [3, 1]> = tensor([[1.0], [0.0], [1.0]]);
val result = x @ y;  // result: Tensor<Float32, [3, 1]>

// Rust — requires ndarray or nalgebra
use ndarray::Array2;
let x = Array2::<f32>::from_shape_vec((3, 3), vec![...]).unwrap();

fn process_batch(input: Tensor<Float32, [_, 784]>): Tensor<Float32, [_, 10]> {
    // batch size (_) is dynamic, feature dimensions are static
    return model.forward(input);
}

// Move tensor to GPU
val gpu_data = data.to(Device.GPU(0));

// GPU-annotated functions
@device(GPU)
fn matmul_kernel(a: Tensor<Float32, [_, _]>, b: Tensor<Float32, [_, _]>): Tensor<Float32, [_, _]> {
    return a @ b;
}

# Rust
cargo build
cargo run
cargo test

# Axon
axonc build main.axon
./main
axonc pkg test

axonc <COMMAND> [OPTIONS]
axonc --help
axonc --version

axonc lex <FILE>

axonc lex hello.axon
# Token(Fn, 1:1)
# Token(Identifier("main"), 1:4)
# Token(LeftParen, 1:8)
# Token(RightParen, 1:9)
# Token(LeftBrace, 1:11)
# ...

axonc parse <FILE> [OPTIONS]