Automatic Differentiation

MIND includes a built-in autodiff engine that generates gradient code at the IR level using reverse-mode automatic differentiation.

Status: shipped, and scoped. Reverse-mode autodiff ships in the open compiler — it is not experimental. It is scoped to the Core v1 tensor ops, entered through a single-output main. Ops outside that set — user functions, control flow, the std surface, modulo, bitwise, shift — return a structured error, never a silently-wrong zero gradient. Neither “experimental” nor “all ops complete” is an accurate summary.
What is spec-level on this page: the executable compiler subset does not yet support first-class function values, so the grad(f)-returns-a-function style shown below, higher-order gradients, custom gradients, and checkpointing are spec-level / roadmap, not shipped executable features. The sections that use them are labelled.

Running it today

Autodiff is feature-gated behind autodiff. Point the compiler at a function and ask for the gradient IR:

mindc program.mind --func main --autodiff --emit-grad-ir

The pipeline is entirely compile-time: it builds a gradient IR mirroring the primal computation, and by default verifies both the primal and gradient modules. The result is deterministic by construction — the gradient IR is built with ordered data structures, so differentiating the same input IR twice produces byte-identical gradient IR and an identical primal-to-gradient value mapping. Only the public IR is touched; no private runtime hooks are referenced.

Basic Usage

Mark functions as differentiable to enable gradient computation:

@differentiable
fn mse_loss(pred: Tensor<f32, N>, target: Tensor<f32, N>) -> f32 {
    mean((pred - target) ** 2)
}

fn main() {
    let pred = [1.0, 2.0, 3.0];
    let target = [1.5, 2.5, 3.5];

    // Compute loss
    let loss = mse_loss(pred, target);

    // Get gradient function
    let grad_fn = grad(mse_loss);
    let d_pred = grad_fn(pred, target);

    print(d_pred);  // Gradient w.r.t. pred
}

How It Works

MIND uses source-transformation reverse-mode AD:

  • Forward pass: Compute output while recording operations
  • Backward pass: Propagate gradients through recorded operations
  • Optimization: Apply standard compiler optimizations to gradient code

Supported Operations

Every differentiable Core v1 operator has a defined gradient rule. A representative sample:

OperationGradient
add(a, b)∂a = upstream, ∂b = upstream
mul(a, b)∂a = upstream * b, ∂b = upstream * a
matmul(a, b)∂a = upstream @ bᵀ, ∂b = aᵀ @ upstream
relu(x)upstream * (x > 0)
sum(x)broadcast(upstream, shape(x))

The full rule set covers:

  • Binary ops — add, sub, mul, div (quotient rule).
  • Matrix ops — dot, matmul (transpose-based rules).
  • Activations — relu, via a dedicated masked-backward op that is Q16.16-preserving.
  • Convolution — conv2d, using dedicated input- and filter-gradient ops. Requires statically-known input and filter shapes; an unknown shape returns an error rather than a guess.
  • Shape-preserving ops — reshape, expand/squeeze dims, slice/index/gather.
  • Reductions — mean (explicit axes) and sum (passthrough).

Fail loud, never fail quiet.Modulo, bitwise, and shift operations are non-differentiable and error explicitly. Two Core v1 operators — tensor.index and tensor.slice— are deliberately excluded from the autodiff contract and surface a clear diagnostic when used in a gradient request. Every unsupported or ambiguous case returns a structured error with a message a caller can act on, so a limitation is reported rather than papered over with a zero gradient.

Higher-Order Gradients (spec-level)

@differentiable
fn f(x: f32) -> f32 {
    x ** 3
}

// First derivative: 3x²
let df = grad(f);

// Second derivative: 6x
let d2f = grad(df);

// Third derivative: 6
let d3f = grad(d2f);

Custom Gradients (spec-level)

@differentiable
@custom_grad(my_relu_grad)
fn my_relu(x: Tensor<f32, N>) -> Tensor<f32, N> {
    max(x, 0.0)
}

fn my_relu_grad(x: Tensor<f32, N>, upstream: Tensor<f32, N>) -> Tensor<f32, N> {
    upstream * cast<f32>(x > 0.0)
}

Gradient Checkpointing (spec-level)

For memory-constrained training, use checkpointing:

@differentiable
@checkpoint  // Recompute forward during backward
fn transformer_block(x: Tensor<f32, B, S, D>) -> Tensor<f32, B, S, D> {
    // Large intermediate activations are not stored
    let attn = self_attention(x);
    let ffn = feed_forward(attn);
    ffn
}

Learn More

See the full autodiff specification at mind-spec/autodiff.md.