Skip to content

refactor: rewrite OverrideInit to be more Zygote compatible #1023

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

Merged
merged 1 commit into from
May 15, 2025

Conversation

AayushSabharwal
Copy link
Member

@AayushSabharwal AayushSabharwal commented May 15, 2025

@DhairyaLGandhi
Copy link
Member

I wouldn't say it's specific to Zygote, but it's also general good practice right? In instances where you cannot determine provenance easily (say you overrode variables with differing types), then also we have to create these copies.

@AayushSabharwal
Copy link
Member Author

Hmm. Sure, I guess. What I find weird is this doesn't work either:

    if initdata.initializeprobmap !== nothing
        u02 = initdata.initializeprobmap(nlsol)
    else
        u02 = u0
    end
    if initdata.initializeprobpmap !== nothing
        p2 = initdata.initializeprobpmap(valp, nlsol)
    else
        p2 = p
    end

    return u02, p2, success

Why is the extra u03 required?

@AayushSabharwal
Copy link
Member Author

Do any other AD frameworks run into this issue? Because if not, it is by definition specific to Zygote regardless of whether this is a code smell or not.

I'm not trying to say Zygote is broken but it would be good to understand this better so we don't keep running into it.

Co-authored-by: Dhairya Gandhi <dhairya@juliahub.com>
@AayushSabharwal AayushSabharwal merged commit 1d06407 into SciML:master May 15, 2025
51 of 60 checks passed
@AayushSabharwal AayushSabharwal deleted the as/make-zygote-happy branch May 15, 2025 19:25
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