Incorrect gradient with `relay.transform.gradient`

I looked into it some further and it seems that ToGraphNormalForm might also be unsound with references… Can you try disabling that also? (Although I’m guessing you want to keep that pass.) You can try using gradient with mode="first_order" if you want to keep the ToGNF pass, since first order will not create references. Hope this helps!