Module autodiff_attrs

This crate handles the user facing autodiff macro. For each #[autodiff(...)] attribute, we create an AutoDiffItem which contains the source and target function names. The source is the function to which the autodiff attribute is applied, and the target is the function getting generated by us (with a name given by the user as the first autodiff arg).

Structs

AutoDiffAttrs

AutoDiffItem

We generate one of these structs for each #[autodiff(...)] attribute.

Enums

DiffActivity

Dual and Duplicated (and their Only variants) are getting lowered to the same Enzyme Activity. However, under forward mode we overwrite the previous shadow value, while for reverse mode we add to the previous shadow value. To not surprise users, we picked different names. Dual numbers is also a quite well known name for forward mode AD types.

DiffMode

Forward and Reverse Mode are well known names for automatic differentiation implementations. Enzyme does support both, but with different semantics, see DiffActivity. The First variants are a hack to support higher order derivatives. We need to compute first order derivatives before we compute second order derivatives, otherwise we would differentiate our placeholder functions. The proper solution is to recognize and resolve this DAG of autodiff invocations, as it’s already done in the C++ and Julia frontend of Enzyme.

Functions

valid_input_activity

valid_ret_activity

Active(Only) is valid in reverse-mode AD for scalar float returns (f16/f32/…). Dual(Only) is valid in forward-mode AD for scalar float returns (f16/f32/…). Const is valid for all cases and means that we don’t compute derivatives wrt. this output. That usually means we have a &mut or *mut T output and compute derivatives wrt. that arg, but this is too complex to verify here. Also it’s just a logic error if users get this wrong.

valid_ty_for_activity

For indirections (ptr/ref) we can’t use Active, since Active allocates a shadow value for the given argument, but we generally can’t know the size of such a type. For scalar types (f16/f32/f64/f128) we can use Active and we can’t use Duplicated, since Duplicated expects a mutable ref/ptr and we would thus end up with a shadow value who is an indirect type, which doesn’t match the primal scalar type. We can’t prevent users here from marking scalars as Duplicated, due to type aliases.