Language

// Function definition (opaque tensor<dtype> accepts any shape)
fn relu(x: tensor<f32>) -> tensor<f32> {
    max(x, 0.0)
}

// Main entry point
fn main() {
    let x: Tensor<f32, 2, 3> = [[1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]];
    let y = relu(x);
    print(y);
}

// Lowercase with commas (Rust-style)
let x: Tensor<f32, 2, 3> = zeros();

// Uppercase with bracket dims (ML-style), concrete dims
let y: Tensor<F32, 4, 8> = zeros();

// Opaque tensor (no shape annotation)
fn relu(x: tensor<f32>) -> tensor<f32> { max(x, 0.0) }

// 1D tensor
let v: Tensor<f32, 3> = [1.0, 2.0, 3.0];

// 2D tensor (matrix)
let m: Tensor<f32, 2, 3> = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];

// Random initialization
let w: Tensor<f32, 784, 128> = randn();

// Zeros/ones
let zeros: Tensor<f32, 10, 10> = zeros();
let ones: Tensor<i32, 5> = ones();

// Regular function with concrete shapes
fn add(a: Tensor<f32, 128>, b: Tensor<f32, 128>) -> Tensor<f32, 128> {
    a + b
}

// Differentiable function (opaque tensor form accepts any shape)
@differentiable
fn loss(pred: tensor<f32>, target: tensor<f32>) -> f32 {
    mean((pred - target) ** 2)
}

// Note: generics are a bounded slice today — a single type parameter over
// scalar types, monomorphized at compile time. Generic shape parameters (N, M)
// have no surface syntax yet; use concrete dims or the opaque tensor<dtype> form

// Comparison operators work on scalars and tensors
let mask = scores > 0.5;        // element-wise comparison
let valid = count >= threshold;  // scalar comparison
let equal = pred == target;      // equality check

// Conditionals
fn activation(x: f32, use_relu: bool) -> f32 {
    if use_relu {
        max(x, 0.0)
    } else {
        tanh(x)
    }
}

// Loops (bounded for determinism)
fn sum_first_n(x: Tensor<f32, 100>, n: i32) -> f32 {
    let mut acc = 0.0;
    for i in 0..n {
        acc = acc + x[i];
    }
    acc
}

// Print for debugging
print(acc);