NaN in gradient when branch not taken has Inf shadow
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
If a branch is not taken in the forward pass, but the shadow of the code in the branch contains `Inf`, the resulting gradient contains `NaN`s. The problem disappears with `Enzyme.API.strong_zero!(true)`, but this should not be needed according to @wsmoses on slack.
```julia-repl
julia> using Enzyme
julia> Enzyme.API.printall!(true)
julia> function f(a, x)
if !isinf(a)
x += a * x
end
return x
end;
julia> autodiff(Reverse, f, Active, Const(Inf), Active(1.0))
after simplification :
; Function Attrs: mustprogress nofree readonly willreturn
define "enzyme_type"="{[-1]:Float@double}" double @preprocess_julia_f_462(double "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4773884592" "enzymejl_parmtype_ref"="0" %0, double "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4773884592" "enzymejl_parmtype_ref"="0" %1) local_unnamed_addr #4 !dbg !22 {
top:
%2 = call {}*** @julia.get_pgcstack() #5
%ptls_field3 = getelementptr inbounds {}**, {}*** %2, i64 2
%3 = bitcast {}*** %ptls_field3 to i64***
%ptls_load45 = load i64**, i64*** %3, align 8, !tbaa !8
%4 = getelementptr inbounds i64*, i64** %ptls_load45, i64 2
%safepoint = load i64*, i64** %4, align 8, !tbaa !12
fence syncscope("singlethread") seq_cst
call void @julia.safepoint(i64* %safepoint) #5, !dbg !23
fence syncscope("singlethread") seq_cst
%5 = call double @llvm.fabs.f64(double %0) #5, !dbg !24
%6 = bitcast double %5 to i64, !dbg !25
%7 = icmp eq i64 %6, 9218868437227405312, !dbg !25
%8 = fmul double %0, %1, !dbg !26
%9 = select i1 %7, double -0.000000e+00, double %8, !dbg !26
%value_phi = fadd double %9, %1, !dbg !26
ret double %value_phi, !dbg !27
}
; Function Attrs: mustprogress nofree
define internal { double } @diffejulia_f_462(double "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4773884592" "enzymejl_parmtype_ref"="0" %0, double "enzyme_type"="{[-1]:Float@double}" "enzymejl_parmtype"="4773884592" "enzymejl_parmtype_ref"="0" %1, double %differeturn) local_unnamed_addr #5 !dbg !28 {
top:
%"value_phi'de" = alloca double, align 8
%2 = getelementptr double, double* %"value_phi'de", i64 0
store double 0.000000e+00, double* %2, align 8
%"'de" = alloca double, align 8
%3 = getelementptr double, double* %"'de", i64 0
store double 0.000000e+00, double* %3, align 8
%"'de1" = alloca double, align 8
%4 = getelementptr double, double* %"'de1", i64 0
store double 0.000000e+00, double* %4, align 8
%"'de2" = alloca double, align 8
%5 = getelementptr double, double* %"'de2", i64 0
store double 0.000000e+00, double* %5, align 8
%6 = call {}*** @julia.get_pgcstack() #6
%ptls_field3 = getelementptr inbounds {}**, {}*** %6, i64 2
%7 = bitcast {}*** %ptls_field3 to i64***
%ptls_load45 = load i64**, i64*** %7, align 8, !tbaa !8, !alias.scope !29, !noalias !32
%8 = getelementptr inbounds i64*, i64** %ptls_load45, i64 2
%safepoint = load i64*, i64** %8, align 8, !tbaa !12, !alias.scope !34, !noalias !37
fence syncscope("singlethread") seq_cst
call void @julia.safepoint(i64* %safepoint) #6, !dbg !39
fence syncscope("singlethread") seq_cst
%9 = call double @llvm.fabs.f64(double %0) #6, !dbg !40
%10 = bitcast double %9 to i64, !dbg !41
%11 = icmp eq i64 %10, 9218868437227405312, !dbg !41
br label %inverttop, !dbg !43
inverttop: ; preds = %top
store double %differeturn, double* %"value_phi'de", align 8
%12 = load double, double* %"value_phi'de", align 8, !dbg !42
store double 0.000000e+00, double* %"value_phi'de", align 8, !dbg !42
%13 = load double, double* %"'de", align 8, !dbg !42
%14 = fadd fast double %13, %12, !dbg !42
store double %14, double* %"'de", align 8, !dbg !42
%15 = load double, double* %"'de1", align 8, !dbg !42
%16 = fadd fast double %15, %12, !dbg !42
store double %16, double* %"'de1", align 8, !dbg !42
%17 = load double, double* %"'de", align 8, !dbg !42
%diffe = select fast i1 %11, double 0.000000e+00, double %17, !dbg !42
store double 0.000000e+00, double* %"'de", align 8, !dbg !42
%18 = load double, double* %"'de2", align 8, !dbg !42
%19 = fadd fast double %18, %17, !dbg !42
%20 = select fast i1 %11, double %18, double %19, !dbg !42
store double %20, double* %"'de2", align 8, !dbg !42
%21 = load double, double* %"'de2", align 8, !dbg !42
store double 0.000000e+00, double* %"'de2", align 8, !dbg !42
%22 = fmul fast double %21, %0, !dbg !42
%23 = load double, double* %"'de1", align 8, !dbg !42
%24 = fadd fast double %23, %22, !dbg !42
store double %24, double* %"'de1", align 8, !dbg !42
fence syncscope("singlethread") seq_cst
fence syncscope("singlethread") seq_cst
%25 = load double, double* %"'de1", align 8
%26 = insertvalue { double } undef, double %25, 0
ret { double } %26
}
((nothing, NaN),)
```
Contributor guide
Assessment
This issue has not been assessed yet.