pub fn graph<F>(name: impl Into<String>, build: F) -> GraphExpand description
Build a single-output graph from a closure.
use rlx_tensor::{graph, shape};
let g = graph("mlp", |g| {
let x = g.input("x", shape![2, 4]);
let w = g.param("w", shape![4, 3]);
let b = g.param("b", shape![3]);
(&x.matmul(&w) + &b).gelu() * 2.0
});
assert_eq!(g.outputs.len(), 1);