Automatic Differentiation
Lux is not an AD package, but it composes well with most of the AD packages available in the Julia ecosystem. This document lists the current level of support for various AD packages in Lux. Additionally, we provide some convenience functions for working with AD.
Overview
AD Package | Mode | CPU | GPU | Nested 2nd Order AD | Support Class |
---|---|---|---|---|---|
ChainRules.jl [1] | Reverse | ✔️ | ✔️ | ✔️ | Tier I |
Enzyme.jl | Reverse | ✔️ | ❓[2] | ❓[2:1] | Tier I[3] |
Zygote.jl | Reverse | ✔️ | ✔️ | ✔️ | Tier I |
ForwardDiff.jl | Forward | ✔️ | ✔️ | ✔️ | Tier I |
ReverseDiff.jl | Reverse | ✔️ | ❌ | ❌ | Tier II |
Tracker.jl | Reverse | ✔️ | ✔️ | ❌ | Tier II |
Mooncake.jl | Reverse | ❓[2:2] | ❌ | ❌ | Tier III |
Diffractor.jl | Forward | ❓[2:3] | ❓[2:4] | ❓[2:5] | Tier III |
Recommendations
For CPU Usacases:
Use
Zygote.jl
for the best performance. This is the most reliable and fastest option for CPU for the time-being. (We are working on faster Enzyme support for CPU)Use
Enzyme.jl
, if there are mutations in the code and/orZygote.jl
fails.If
Enzyme.jl
fails for some reason, (open an issue and) tryReverseDiff.jl
(possibly with compiled mode).
For GPU Usacases:
- Use
Zygote.jl
for the best performance. This is the most reliable and fastest option for GPU for the time-being. We are working on supportingEnzyme.jl
for GPU as well.
- Use
Support Class
Tier I: These packages are fully supported and have been tested extensively. Often have special rules to enhance performance. Issues for these backends take the highest priority.
Tier II: These packages are supported and extensively tested but often don't have the best performance. Issues against these backends are less critical, but we fix them when possible. (Some specific edge cases, especially with AMDGPU, are known to fail here)
Tier III: We don't know if these packages currently work with Lux. We'd love to add tests for these backends, but currently these are not our priority.
Footnotes
Note that
ChainRules.jl
is not really an AD package, but we have first-class support for packages that userrules
. ↩︎This feature is supported downstream, but we don't extensively test it to ensure that it works with Lux. ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎
Currently Enzyme outperforms other AD packages in terms of CPU performance. However, there are some edge cases where it might not work with Lux. We are working on improving the compatibility. Please report any issues you encounter. ↩︎