Skip to main content

Crate ferrotorch_jit_script

Crate ferrotorch_jit_script 

Source
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 in ferrotorch_core::grad_fns::{arithmetic, reduction, activation, linalg} that takes &Tensor args and returns Tensor)
  • 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 function fn(args...) -> Tensor<T> to compile it into a ferrotorch_jit::TracedModule<T>-returning function.