RLX tab cluster — one rank per tab

import init, {
  WebGraph, compile_forward, compile_backward,
} from "./pkg/rlx_web.js";

await init();                     // load the wasm module

// ── describe the model ──
// 784 → 128 → 10 with softmax cross-entropy. Node ids are
// plain numbers; `param` marks a tensor as trainable.
const BATCH = 64;
const g = new WebGraph("mlp");
const x      = g.input("x", [BATCH, 784]);
const labels = g.input("labels", [BATCH]);  // class idx/row
const w1 = g.param("w1", [784, 128]);
const b1 = g.param("b1", [128]);
const w2 = g.param("w2", [128, 10]);
const b2 = g.param("b2", [10]);

let h = g.matmul(x, w1, [BATCH, 128]);
h = g.add(h, b1, [BATCH, 128]);
h = g.relu(h, [BATCH, 128]);
let logits = g.matmul(h, w2, [BATCH, 10]);
logits = g.add(logits, b2, [BATCH, 10]);

const perSample = g.softmax_cross_entropy(logits, labels);
g.set_outputs([g.mean(perSample, 0, [])]);     // scalar loss

// ── call it ──
// "cpu" → "webgpu", "webgl" or "auto"; same graph.
const fwd = compile_forward(g, "cpu");
fwd.set_param("w1", W1); fwd.set_param("b1", B1);
fwd.set_param("w2", W2); fwd.set_param("b2", B2);
fwd.feed("x", xs);
fwd.feed("labels", ys);
const [loss] = fwd.run();         // output_sizes() is [1]

// ── train it ──
// Autodiff rewrites the graph: outputs become
// [loss, dW1, db1, dW2, db2], and a `d_output` seed input
// appears for ∂loss/∂loss.
const step = compile_backward(g, [w1, b1, w2, b2], "cpu");
const sizes = step.output_sizes();
const weights = [W1, B1, W2, B2];
const NAMES = ["w1", "b1", "w2", "b2"];

step.feed("x", xs);
step.feed("labels", ys);
step.feed("d_output", new Float32Array([1]));

for (let i = 0; i < 100; i++) {
  NAMES.forEach((n, k) => step.set_param(n, weights[k]));
  const out = step.run();           // [loss, grads…] flat
  let at = sizes[0];
  weights.forEach((w, k) => {
    const n = sizes[k + 1];
    for (let j = 0; j < n; j++) w[j] -= 0.05 * out[at + j];
    at += sizes[k + 1];
  });
  if (i % 20 === 0) console.log(i, out[0]);
}

// On WebGPU readback cannot block the loop — await it:
//   const out = await step.run_async();

use rlx_autodiff::grad_with_loss;
use rlx_ir::{op::*, DType, Graph, NodeId, Shape};
use rlx_optim::{Optimizer, Sgd};
use rlx_runtime::{Device, Session};

const BATCH: usize = 64;

/// 784 → 10 classifier with softmax cross-entropy. Returns
/// the graph and the parameters to differentiate wrt.
fn build_loss() -> (Graph, Vec<NodeId>) {
    let f = DType::F32;
    let mut g = Graph::new("mlp");
    let x = g.input("x", Shape::new(&[BATCH, 784], f));
    let lbl = g.input("labels", Shape::new(&[BATCH], f));
    let w = g.param("w", Shape::new(&[784, 10], f));
    let b = g.param("b", Shape::new(&[10], f));

    let mm = g.matmul(x, w, Shape::new(&[BATCH, 10], f));
    let shape = Shape::new(&[BATCH, 10], f);
    let logits = g.binary(BinaryOp::Add, mm, b, shape);
    let per = g.softmax_cross_entropy_with_logits(logits, lbl);
    let loss = g.add_node(
        Op::Reduce {
            op: ReduceOp::Mean,
            axes: vec![0],
            keep_dim: false,
        },
        vec![per],
        Shape::from_dims(&[], f),
    );
    g.set_outputs(vec![loss]);
    (g, vec![w, b])
}

// ── call it ──
// `Device::Cpu` → `Metal`, `Cuda`, `WebGpu`, … same graph.
let (graph, params) = build_loss();
let mut fwd = Session::new(Device::Cpu).compile(graph.clone());
fwd.set_param("w", &w_data);
fwd.set_param("b", &b_data);
let feed = [("x", &x_data[..]), ("labels", &labels[..])];
let loss = fwd.run(&feed)[0][0];

// ── train it ──
// Autodiff rewrites the graph; outputs become [loss, dW, db].
let backward = grad_with_loss(&graph, &params);
let mut step = Session::new(Device::Cpu).compile(backward);
let mut sgd = Sgd::new(0.05);

for _ in 0..100 {
    step.set_param("w", &w_data);
    step.set_param("b", &b_data);
    let outs = step.run(&[
        ("x", &x_data),
        ("labels", &labels),
        ("d_output", &[1.0][..]),   // seed ∂loss/∂loss = 1
    ]);
    sgd.step("w", &[784, 10], &mut w_data, &outs[1]);
    sgd.step("b", &[10], &mut b_data, &outs[2]);
}

import numpy as np
import pyrlx as rlx

device = rlx.available_devices()[0]   # "cpu", "metal", …

with rlx.graph("mlp") as g:
    x = g.input("x", [64, 784], "f32")
    target = g.input("target", [64, 10], "f32")
    w = g.param("w", [784, 10], "f32")
    b = g.param("b", [10], "f32")

    y = (x @ w + b).relu()
    loss = ((y - target) ** 2).mean([0, 1])    # MSE
    g.outputs = [loss]
    forward, w_id, b_id = g.raw, w.id, b.id

W = (np.random.randn(784, 10) * 0.02).astype(np.float32)
B = np.zeros(10, dtype=np.float32)
xs = np.random.randn(64, 784).astype(np.float32)
ts = np.zeros((64, 10), dtype=np.float32)

# ── call it ──
compiled = rlx.Session(device).compile(forward)
rlx.set_param(compiled, "w", W)
rlx.set_param(compiled, "b", B)
(loss,) = rlx.run(compiled, x=xs, target=ts)
print("loss:", float(loss))

# ── train it ──
# `grad` returns a graph whose outputs are [loss, dW, db].
with rlx.graph("mlp") as g2:          # `g.raw` consumed g
    x = g2.input("x", [64, 784], "f32")
    target = g2.input("target", [64, 10], "f32")
    w = g2.param("w", [784, 10], "f32")
    b = g2.param("b", [10], "f32")
    y = (x @ w + b).relu()
    g2.outputs = [((y - target) ** 2).mean([0, 1])]
    fwd2, ids = g2.raw, [w.id, b.id]

backward = rlx.grad(fwd2, ids)
step = rlx.Session(device).compile(backward)
seed = np.ones(1, dtype=np.float32)            # ∂loss/∂loss

for i in range(100):
    rlx.set_param(step, "w", W)
    rlx.set_param(step, "b", B)
    loss, dW, dB = rlx.run(step, x=xs, target=ts, d_output=seed)
    W -= 0.05 * dW.reshape(W.shape)
    B -= 0.05 * dB.reshape(B.shape)
    if i % 20 == 0:
        print(i, float(loss))