Trait TracedTensorAdExt

Source

pub trait TracedTensorAdExt {
    // Required methods
    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor>;
    fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>>;
    fn checkpoint(
        &mut self,
        compiler: &mut GraphCompiler,
        runtime: &Runtime,
    ) -> Result<()>;
    fn jvp(
        &self,
        wrt: &TracedTensor,
        tangent: &TracedTensor,
    ) -> Result<TracedTensor>;
    fn jvp_optional(
        &self,
        wrt: &TracedTensor,
        tangent: &TracedTensor,
    ) -> Result<Option<TracedTensor>>;
    fn vjp(
        &self,
        wrt: &TracedTensor,
        cotangent: &TracedTensor,
    ) -> Result<TracedTensor>;
    fn vjp_optional(
        &self,
        wrt: &TracedTensor,
        cotangent: &TracedTensor,
    ) -> Result<Option<TracedTensor>>;
}

Expand description

Automatic differentiation helpers for [TracedTensor].

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_runtime::TracedTensor;

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let loss = x.scale_real(2.0).unwrap();
let maybe_dx = loss.grad_optional(&x).unwrap();
assert!(maybe_dx.is_some());

Required Methods§

Source

fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor>

Gradient of a scalar output with respect to a traced input.

For complex scalar outputs, tenferro returns the Hermitian-adjoint cotangent. To compare seed-1 scalar gradients with JAX’s public grad values, use the complex conjugate of this result. See https://tensor4all.org/tenferro-rs/guides/complex-ad.html.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuBackend;
use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};

fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
    let mut compiler = GraphCompiler::new();
    let program = compiler.compile(tensor).unwrap();
    let backend = CpuBackend::new();
    let mut builder = Runtime::builder();
    builder
        .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
        .unwrap();
    let runtime = builder.build().unwrap();
    runtime.run_compiled(&program, &[]).unwrap().pop().unwrap()
}

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let loss = (&x * &x).unwrap();
let dx = loss.grad(&x).unwrap();

assert_eq!(eval(&dx).as_slice::<f64>().unwrap(), &[6.0]);

§Errors

Returns tenferro_runtime::Error::NonScalarGrad for a non-scalar output, tenferro_runtime::Error::UnsupportedAdRule when an AD rule is unavailable, or a typed validation/backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Source

fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>>

Like grad, but returns None when wrt is inactive.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_runtime::TracedTensor;

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
let loss = (&y * &y).unwrap();

assert!(loss.grad_optional(&x).unwrap().is_none());

§Errors

Returns tenferro_runtime::Error::NonScalarGrad for a non-scalar output, Error::UnsupportedAdRule when an AD rule is unavailable, or a typed validation/backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Source

fn checkpoint( &mut self, compiler: &mut GraphCompiler, runtime: &Runtime, ) -> Result<()>

Evaluate this tensor and replace its graph with a concrete leaf while preserving the previous graph for AD replay.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuBackend;
use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};

let mut compiler = GraphCompiler::new();
let backend = CpuBackend::new();
let mut builder = Runtime::builder();
builder
    .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
    .unwrap();
let runtime = builder.build().unwrap();
let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let mut y = (&x * &x).unwrap();

y.checkpoint(&mut compiler, &runtime).unwrap();

let value = y.attached_data().unwrap();
assert_eq!(value.as_slice::<f64>().unwrap(), &[9.0]);

§Errors

Returns tenferro_runtime::Error::Validation when checkpoint metadata is invalid, Error::RuntimeState when graph metadata or runtime state is unavailable, or a typed backend error from evaluation.

Source

fn jvp( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<TracedTensor>

Forward-mode Jacobian-vector product.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuBackend;
use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};

fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
    let mut compiler = GraphCompiler::new();
    let program = compiler.compile(tensor).unwrap();
    let backend = CpuBackend::new();
    let mut builder = Runtime::builder();
    builder
        .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
        .unwrap();
    let runtime = builder.build().unwrap();
    runtime.run_compiled(&program, &[]).unwrap().pop().unwrap()
}

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let tangent = TracedTensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap();
let y = (&x * &x).unwrap();
let dy = y.jvp(&x, &tangent).unwrap();

assert_eq!(eval(&dy).as_slice::<f64>().unwrap(), &[12.0]);

§Errors

Returns tenferro_runtime::Error::UnsupportedAdRule when a JVP rule is unavailable, Error::Validation for incompatible tangent metadata, or a typed backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Source

fn jvp_optional( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

Like jvp, but returns None when wrt is inactive.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_runtime::TracedTensor;

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
let tangent = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
let loss = (&y * &y).unwrap();

assert!(loss.jvp_optional(&x, &tangent).unwrap().is_none());

§Errors

Returns tenferro_runtime::Error::UnsupportedAdRule when a JVP rule is unavailable, Error::Validation for incompatible tangent metadata, or a typed backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Source

fn vjp( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<TracedTensor>

Reverse-mode vector-Jacobian product.

Complex cotangents use tenferro’s Hermitian real-inner-product convention. Non-real complex cotangent seeds therefore need an explicit seed-convention comparison when matching JAX. See https://tensor4all.org/tenferro-rs/guides/complex-ad.html.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuBackend;
use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};

fn eval(tensor: &TracedTensor) -> tenferro_runtime::Tensor {
    let mut compiler = GraphCompiler::new();
    let program = compiler.compile(tensor).unwrap();
    let backend = CpuBackend::new();
    let mut builder = Runtime::builder();
    builder
        .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
        .unwrap();
    let runtime = builder.build().unwrap();
    runtime.run_compiled(&program, &[]).unwrap().pop().unwrap()
}

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let cotangent = TracedTensor::from_vec_col_major(vec![], vec![0.5_f64]).unwrap();
let y = (&x * &x).unwrap();
let dx = y.vjp(&x, &cotangent).unwrap();

assert_eq!(eval(&dx).as_slice::<f64>().unwrap(), &[3.0]);

§Errors

Returns tenferro_runtime::Error::UnsupportedAdRule when a VJP rule is unavailable, Error::Validation for incompatible cotangent metadata, or a typed backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Source

fn vjp_optional( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

Like vjp, but returns None when wrt is inactive.

§Examples

use tenferro_ad::TracedTensorAdExt;
use tenferro_runtime::TracedTensor;

let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
let y = TracedTensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap();
let cotangent = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
let loss = (&y * &y).unwrap();

assert!(loss.vjp_optional(&x, &cotangent).unwrap().is_none());

§Errors

Returns tenferro_runtime::Error::UnsupportedAdRule when a VJP rule is unavailable, Error::Validation for incompatible cotangent metadata, or a typed backend/runtime-state error.

§Deferred errors

Symbolic shape constraints can later produce tenferro_runtime::Error::ShapeConstraintViolation or tenferro_runtime::Error::ShapeConstraintEvaluation during compile or execution.

Implementations on Foreign Types§

Source §

fn checkpoint( &mut self, compiler: &mut GraphCompiler, runtime: &Runtime, ) -> Result<()>

Source §

fn jvp( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<TracedTensor>

Source §

fn jvp_optional( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

Source §

fn vjp( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<TracedTensor>

Source §

TracedTensorAdExt

Trait TracedTensorAdExt Copy item path

§Examples

Required Methods§

fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor>

§Examples

§Errors

§Deferred errors

fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>>

§Examples

§Errors

§Deferred errors

fn checkpoint( &mut self, compiler: &mut GraphCompiler, runtime: &Runtime, ) -> Result<()>

§Examples

§Errors

fn jvp( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<TracedTensor>

§Examples

§Errors

§Deferred errors

fn jvp_optional( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

§Examples

§Errors

§Deferred errors

fn vjp( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<TracedTensor>

§Examples

§Errors

§Deferred errors

fn vjp_optional( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

§Examples

§Errors

§Deferred errors

Implementations on Foreign Types§

impl TracedTensorAdExt for TracedTensor

fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor>

fn grad_optional(&self, wrt: &TracedTensor) -> Result<Option<TracedTensor>>

fn checkpoint( &mut self, compiler: &mut GraphCompiler, runtime: &Runtime, ) -> Result<()>

fn jvp( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<TracedTensor>

fn jvp_optional( &self, wrt: &TracedTensor, tangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

fn vjp( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<TracedTensor>

fn vjp_optional( &self, wrt: &TracedTensor, cotangent: &TracedTensor, ) -> Result<Option<TracedTensor>>

Implementors§

Trait TracedTensorAdExt