ChainRules.jl meets Unitful.jl: Autodiff via Unit Analysis

Tools for performing autodifferentiation (AD) and dimensional work in Julia are robust, but not always compatible. This talk explores how we can understand rule-based AD in Julia by showing how to make dimensional quantities from Unitful.jl compose with ChainRules.jl. Combining these two projects produces an intuitive look at the building blocks of AD in Julia using only rudimentary calculus and dimensional analysis.

Unitful.jl provides efficient type-level support for dimensional quantities we encounter when simulating physical systems. Likewise, ChainRules.jl forms the backbone of robust but easily-extensible autodifferentiation (AD) systems. Exploring these two systems together yields an insightful look at Julia's rule-based AD. Calculus, dimensional analysis, and physical intuition are sufficient to explain how ChainRules.jl works by building AD rules for Unitful.jl.

The versatility of the ChainRules.jl ecosystem arises from implementing and extending a ruleset for fundamental functions, such as * and inv, as rrules and frules. What can often seem like a mysterious black box that computes derivatives is actually composed of many individual rrules or frules built on rudimentary calculus.

These rrules and frules are interpreted by thinking about differentiation as a problem of physical dimensions, and Unitful.jl is used to confirm these findings. However, arithmetic between Unitful.jl quantities is not immediately compatible with ChainRules.jl-based AD. This talk presents the pertinent AD rules to enable basic ChainRules.jl compatibility. These rules are also used as a lens to understand how to read and write AD rules for the ChainRules.jl ecosystem.