make_zero! for immutable

[vchuravy]

Recurse through the immutable and set the mutable fields to zero.

[danielwe]

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!: danielwe@7a6ca9f. 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.

[vchuravy]

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.

[danielwe]

If prev is an immutable the line

setfield!(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: danielwe@7a6ca9f#diff-10924b07e20c4a235afd58dea795108f564d5f535318adb73991e4e57144e09dR211-R223

Is this what you've been running into?

[danielwe]

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.

[danielwe]

...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.

[wsmoses]

@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)

[danielwe]

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)

[wsmoses]

okay, I'll close this for now then (and definitely reopen if needed)

[vchuravy]

We definilty need one of the PRs that @danielwe is working on.

[danielwe]

#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.