autodiff_forward

Attribute Macro autodiff_forward 

Source
#[autodiff_forward]
🔬This is a nightly-only experimental API. (autodiff #124509)
Expand description

This macro handles automatic differentiation. This macro uses forward-mode automatic differentiation to generate a new function. It may only be applied to a function. The new function will compute the derivative of the function to which the macro was applied.

The expected usage syntax is: #[autodiff_forward(NAME, INPUT_ACTIVITIES, OUTPUT_ACTIVITY)]

  • NAME: A string that represents a valid function name.
  • INPUT_ACTIVITIES: Specifies one valid activity for each input parameter.
  • OUTPUT_ACTIVITY: Must not be set if the function implicitly returns nothing (or explicitly returns -> ()). Otherwise, it must be set to one of the allowed activities.

ACTIVITIES might either be Dual or Const, more options will be exposed later.

Const should be used on non-float arguments, or float-based arguments as an optimization if we are not interested in computing the derivatives with respect to this argument.

Dual can be used for float scalar values or for references, raw pointers, or other indirect input arguments. It can also be used on a scalar float return value. If used on a return value, the generated function will return a tuple of two float scalars. If used on an input argument, a new shadow argument of the same type will be created, directly following the original argument.

§Usage examples:

#![feature(autodiff)]
use std::autodiff::*;
#[autodiff_forward(rb_fwd1, Dual, Const, Dual)]
#[autodiff_forward(rb_fwd2, Const, Dual, Dual)]
#[autodiff_forward(rb_fwd3, Dual, Dual, Dual)]
fn rosenbrock(x: f64, y: f64) -> f64 {
    (1.0 - x).powi(2) + 100.0 * (y - x.powi(2)).powi(2)
}
#[autodiff_forward(rb_inp_fwd, Dual, Dual, Dual)]
fn rosenbrock_inp(x: f64, y: f64, out: &mut f64) {
    *out = (1.0 - x).powi(2) + 100.0 * (y - x.powi(2)).powi(2);
}

fn main() {
  let x0 = rosenbrock(1.0, 3.0); // 400.0
  let (x1, dx1) = rb_fwd1(1.0, 1.0, 3.0); // (400.0, -800.0)
  let (x2, dy1) = rb_fwd2(1.0, 3.0, 1.0); // (400.0, 400.0)
  // When seeding both arguments at once the tangent return is the sum of both.
  let (x3, dxy) = rb_fwd3(1.0, 1.0, 3.0, 1.0); // (400.0, -400.0)

  let mut out = 0.0;
  let mut dout = 0.0;
  rb_inp_fwd(1.0, 1.0, 3.0, 1.0, &mut out, &mut dout);
  // (out, dout) == (400.0, -400.0)
}

We might want to track how one input float affects one or more output floats. In this case, the shadow of one input should be initialized to 1.0, while the shadows of the other inputs should be initialized to 0.0. The shadow of the output(s) should be initialized to 0.0. After calling the generated function, the shadow of the input will be zeroed, while the shadow(s) of the output(s) will contain the derivatives. Forward mode is generally more efficient if we have more output floats marked as Dual than input floats. Related information can also be found under the term “Vector-Jacobian product” (VJP).