Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Just call primal on intrinics if all tangents are zero #282

Merged
merged 1 commit into from
Mar 18, 2024
Merged

Conversation

oxinabox
Copy link
Member

Encountered this in the wild.
I can't however reproduce this in a simple test like

@testset "zeroed tangents intrinstics" begin
    foo(x,y) = Core.Intrinsics.and_int(x, y)
    z = ∂☆internal{1}()(DNEBundle{1}(foo), DNEBundle{1}(true), DNEBundle{1}(false))
    @test primal(z) == foo(true, false)
    @test first_partial
end

Not sure why. Maybe the more general handling of all DNEBundle applies there but for some reason in my real world example?

An alternative to this is to change our general fallback handling of intrinsics to handle zeros, whether compiletime known or not

function (this::∂☆{N})(f::AbstractZeroBundle{N, Core.IntrinsicFunction}, args::ATB{N}...) where {N}
    ff = primal(f)
    primal_args = map(primal, args)
    primal_out = ff(primal_args...)
    if all(iszero, primal_args))
        (zero_bundle{N}())(primal_out)
    elseif ff in (Base.not_int, Base.ne_float)
        DNEBundle{N}(primal_out)
    else
        error("Missing rule for intrinsic function $ff")
    end
end

@Keno
Copy link
Collaborator

Keno commented Mar 18, 2024

I've generally been resistant to making this part of the base Diffractor semantics, because people have occasionally written rules that violate this, e.g. the Flux rules for testing whether something is happening under differentiation or the batchnorm update rules. I also don't think it's fully necessary, because stage 2 diffractor should be able to analyze the code and determine whether this is true. That said, if this fixes a short term issue, I'm fine with it for now.

@oxinabox
Copy link
Member Author

We have this for things that are not intrinsics and have had for a long time.
Practically it is very useful for avoiding ADing code not on the main path.
This particular case for incidence actually seems less objectionable on that reason.

But yes, moving to traving what is on the active path directly is my preference, longer term.
Not so much for things like batch norm update rules though, I kinda think those shouldn't be done via writing frules
and rather should be using a ScopedValue, it feels better to me.

@oxinabox oxinabox merged commit 0f74c78 into main Mar 18, 2024
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants