Skip to main content

script

Attribute Macro script 

Source
#[script]
Expand description

Apply #[script] to a function fn(args...) -> Tensor<T> to compile it into a ferrotorch_jit::TracedModule<T>-returning function.

Example:

use ferrotorch_jit_script::script;
use ferrotorch_core::Tensor;
use ferrotorch_core::grad_fns::arithmetic::{mul, add};
use ferrotorch_core::grad_fns::reduction::sum;

#[script]
fn weighted_sum(a: Tensor<f32>, w: Tensor<f32>) -> Tensor<f32> {
    let prod = mul(&a, &w)?;
    sum(&prod)
}

// `weighted_sum(...)` now returns `FerrotorchResult<TracedModule<f32>>`
// built by tracing the body once with the supplied tensors.