Skip to main content

graph

Function graph 

Source
pub fn graph<F>(name: impl Into<String>, build: F) -> Graph
where F: FnOnce(&mut GraphScope) -> Tensor,
Expand 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);