Document CheckNaN

EnzymeAD · Feb 16, 2025 · 0b4fee5 · 0b4fee5
1 parent 4a5316e
commit 0b4fee5
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -3,7 +3,7 @@
 ## Types and constants
 
 ```@index
-Modules = [Enzyme, EnzymeCore, EnzymeCore.EnzymeRules, EnzymeTestUtils, Enzyme.API]
+Modules = [Enzyme, EnzymeCore, EnzymeCore.EnzymeRules, EnzymeTestUtils, Enzyme.API, Enzyme.Compiler]
 Order = [:type, :constant]
 ```
 
@@ -20,3 +20,7 @@ Order = [:macro, :function]
 Modules = [Enzyme, EnzymeCore, EnzymeCore.EnzymeRules, EnzymeTestUtils, Enzyme.API]
 Order = [:module, :type, :constant, :macro, :function]
 ```
+
+```@docs
+Enzyme.Compiler.CheckNan
+```
diff --git a/docs/src/internal_api.md b/docs/src/internal_api.md
@@ -9,4 +9,5 @@
 ```@autodocs
 Modules = [Enzyme.Compiler]
 Order = [:module, :type, :constant, :macro, :function]
+Filter = t -> !(t === Enzyme.Compiler.CheckNan)
 ```
diff --git a/src/compiler.jl b/src/compiler.jl
@@ -437,7 +437,15 @@ function removed_ret_parms(F::LLVM.Function)
     return retRemove, parmsRemoved
 end
 
+"""
+    CheckNan::Ref{Bool}
+
+If `Enzyme.Compiler.CheckNan[] == true`, Enzyme will error at the first encounter of a `NaN`
+during differentiation. Useful as a debugging tool to help locate the call whose derivative
+is the source of unexpected `NaN`s. Off by default.
+"""
 const CheckNan = Ref(false)
+
 function julia_sanitize(
     orig::LLVM.API.LLVMValueRef,
     val::LLVM.API.LLVMValueRef,