Equivalent functions but different gradients

Detailed description in issue 70578


However, in the attempt of the following code, it seems that the correct gradient is obtained.

I wonder if anyone can help me solve this problem or give me some debugging suggestions

code