Crate chainrules 

AD engine: tape-based reverse-mode and dual-number forward-mode.

This crate provides the AD execution engine, built on top of [chainrules_core] traits. It is analogous to Zygote.jl in the Julia ecosystem: a concrete AD engine that uses ChainRulesCore.jl interfaces.

• Reverse-mode AD via Tape, TrackedTensor, and Tape::pullback
• Forward-mode AD via DualTensor
• Forward-over-reverse HVP via Tape::hvp

Operation-specific AD rules (e.g., einsum rrule/frule) live in the crate that defines the operation. See tenferro-einsum for einsum AD functions.

Bodies are intentionally todo!() in the current POC phase.

Examples

Reverse-mode usage (with operation-specific AD functions from other crates):

use chainrules::{Tape, TrackedTensor};
use tenferro_einsum::tracked_einsum;
use tenferro_tensor::{MemoryOrder, Tensor};
use tenferro_device::LogicalMemorySpace;

let tape = Tape::<Tensor<f64>>::new();
let a = tape.leaf(Tensor::ones(
    &[2, 3],
    LogicalMemorySpace::MainMemory,
    MemoryOrder::ColumnMajor,
));
let b = tape.leaf(Tensor::ones(
    &[3, 4],
    LogicalMemorySpace::MainMemory,
    MemoryOrder::ColumnMajor,
));
let c = tracked_einsum("ij,jk->ik", &[&a, &b]).unwrap();
let loss = tracked_einsum("ij,ij->", &[&c, &c]).unwrap();
let grads = tape.pullback(&loss).unwrap();
let _ga = grads.get(a.node_id().unwrap()).unwrap();

Forward-mode usage:

use chainrules::DualTensor;
use tenferro_einsum::dual_einsum;
use tenferro_tensor::{MemoryOrder, Tensor};

let a = Tensor::<f64>::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2], MemoryOrder::ColumnMajor).unwrap();
let da = Tensor::<f64>::ones(&[2, 2], tenferro_device::LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor);
let b = Tensor::<f64>::ones(&[2, 2], tenferro_device::LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor);

let a_dual = DualTensor::with_tangent(a, da).unwrap();
let b_dual = DualTensor::new(b);
let c_dual = dual_einsum("ij,jk->ik", &[&a_dual, &b_dual]).unwrap();
let _jvp = c_dual.tangent();

Forward-over-reverse HVP (Hessian-vector product):

use chainrules::Tape;
use tenferro_einsum::tracked_einsum;
use tenferro_tensor::{MemoryOrder, Tensor};
use tenferro_device::LogicalMemorySpace;

let tape = Tape::<Tensor<f64>>::new();
let x = tape.leaf_with_tangent(
    Tensor::ones(&[3], LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor),
    Tensor::ones(&[3], LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor),  // direction v
).unwrap();
let loss = tracked_einsum("i,i->", &[&x, &x]).unwrap();  // f(x) = x·x
let result = tape.hvp(&loss).unwrap();
let _grad = result.gradients;  // ∇f(x) = 2x
let _hv = result.hvp;          // H·v = 2v

Structs

DualTensor

Value wrapper for forward-mode AD.

Gradients

Accumulated gradients indexed by NodeId.

HvpResult

Result of a forward-over-reverse HVP computation.

NodeId

Stable identifier of an AD graph node.

PullbackPlan

Compiled pullback execution plan.

Tape

Reverse-mode AD tape.

TrackedTensor

Value wrapper for reverse-mode AD.

Enums

AutodiffError

AD-specific error type.

SavePolicy

Saved-tensor retention policy for reverse-mode rules.

Traits

Differentiable

Trait defining the tangent space for a differentiable type.

ForwardRule

Forward-mode AD rule interface (frule).

ReverseRule

Reverse-mode AD rule interface (rrule).

Type Aliases

AdResult

Result alias for AD APIs.