Expand description
#[ferrotorch_jit_script::script] — declarative graph construction. (#625)
Annotate a Rust function with #[script] and the macro rewrites the body
to build an IrGraph instead of running eagerly. The wrapped function
returns a closure with the same signature that, when called, executes
the captured graph via ferrotorch_jit::trace.
§What’s supported in the body
The script body is restricted to a small recognized subset:
let x = …;bindings (no shadowing across statements is required)- Function calls on tensors:
add(&a, &b),mul(&a, &b),sum(&t),mean(&t),relu(&t),sigmoid(&t),tanh(&t),mm(&a, &b)… (anything inferrotorch_core::grad_fns::{arithmetic, reduction, activation, linalg}that takes&Tensorargs and returnsTensor) - A trailing expression as the function’s return value
Anything else (control flow, struct fields, non-tensor values) is left
untouched — the macro just emits the rewritten body as-is and trace
captures it at the autograd level.
§Why this is a shim over trace
The macro doesn’t reimplement the IR builder. Instead, it ensures the
function’s inputs are leaf tensors with requires_grad=true (so the
autograd graph is built), then calls the existing trace to capture
the IR. That keeps op coverage in lockstep with trace and avoids a
second source of truth for op-name → IR mapping.
Attribute Macros§
- script
- Apply
#[script]to a functionfn(args...) -> Tensor<T>to compile it into aferrotorch_jit::TracedModule<T>-returning function.