EnzymeAD
make_zero! for immutable
Hi, just want to mention that I've spent a lot of time with make_zero! lately for #1852, which is almost ready. There are several bugs (missing branches in some make_zero_immutable! methods, several in(seen, prev) that should be in(prev, seen), et cetera), but I'm not sure a separate !ismutable branch in the generic method is needed. These cases should go through line ~404~ 398 (EDIT: changed line # to refer to the pre-PR state). What's the error you're seeing?
You can check out this commit on a branch in my fork for fixes to all the bugs I found in make_zero and make_zero!: https://github.com/danielwe/Enzyme.jl/commit/7a6ca9ff332b413c80a2cc540a92c79eb026f2eb. I haven't PRd it because it's redundant with #1852, but if a stopgap solution is needed, you might find the fix you need there.
What's the error you're seeing?
If prev is an immutable the line
setfield!(prev, i, make_zero_immutable!(xi, seen))
in the generic handler will fail.
If
previs an immutable the linesetfield!(prev, i, make_zero_immutable!(xi, seen))in the generic handler will fail.
That's only the case if the immutable contains immutable differentiable values, like (; a=1.0, b=[1.0]) (ActiveState or MixedState), in which case make_zero! should fail anyway. If make_zero! succeeds on an immutable type, that's because every differentiable value is contained within mutable storage, like (; a=[1.0], b=[1.0]) and the test case in this PR (DupState or, trivially, AnyState). In these cases, you will never hit the setfield! branch, you'll hit the branch below that's recursively calling make_zero!.
I.e., the test case added in this PR would never hit the setfield! branch anyway.
However, if you have a MixedState NamedTuple contained within a mutable container, such as [(; a=1.0, b=[1.0])], and call make_zero! on the outer container, then you'll hit the make_zero_immutable!(::Type{<:NamedTuple}, ...) method, and this one has a bug where it fails to call back into make_zero! for the mutable fields. That's the bug fixed in this part of the commit I linked: https://github.com/danielwe/Enzyme.jl/commit/7a6ca9ff332b413c80a2cc540a92c79eb026f2eb#diff-10924b07e20c4a235afd58dea795108f564d5f535318adb73991e4e57144e09dR211-R223
Is this what you've been running into?
Just realized the proposed fix here could also be a way to mask #1935 (because the bug in active_reg_inner makes the recursion take the wrong branch and call make_zero! when it should have called make_zero_immutable!), but I think the proper fix for that should be #1936 or something similar. The fix proposed here enables silent bugs where make_zero! fails to throw an error on unsupported types, while the bug in active_reg_inner remains unaddressed and may show up in other contexts.
...and any issue addressed by this PR that isn't caused by #1935 should be fixed by #1961. If that's not the case, please send me an MWE so I can fix that.
@vchuravy @danielwe is this PR still live or has it been partially covered elsewhere (please forgive the delay, am getting back to this after the big compile time fix)
Pretty sure the underlying issues here were fixed by #1961. The test suite introduced there should cover every failure mode this PR aimed to address. @vchuravy please let me know if you have an MWE that still fails.
That said, I was a bit too confident in the discussion above. The issue was a bit more nuanced and #1961 didn't get it right until after the full combinatorial nested wrapper testset was added, which was after this discussion. (The solution was to use active_reg_inner less and ismutabletype more systematically throughout the make_zero! methods, as seen for example in this section of the diff: https://github.com/EnzymeAD/Enzyme.jl/pull/1961/files#diff-10924b07e20c4a235afd58dea795108f564d5f535318adb73991e4e57144e09dL441-R500)
#1961 was already merged, so all this should work in the latest releases!
The bigger #1852 will be ready soon(TM), but that one shouldn't affect make_zero(!) correctness now.