Power and multiple multiplication don't give the same gradient

[ASionville]

🐛 Describe the bug

Hello, I have an issue with a code of mine :
I have a function, and I want to compute its derivatives using Fourier Transform.

First I compute the waves numbers :
k_x = tensor([[[ 0, 1, 2, 3, 4, 5, 6, 7, -8, -7, -6, -5, -4, -3, -2, -1]]])

Then I get the first and second derivatives :

ft_dudx = 2j * np.pi * k_x * ft_u
ft_dudx2 = 2j * np.pi * k_x * ft_dudx # = 2j * np.pi * k_x * 2j * np.pi * k_x * ft_u

This method works well, and using the classic optimizer method I can get the correct gradients and lower the loss

But when I do this, which is mathematically equivalent, it doesn't work at all :

ft_dudx = ((2j * np.pi * k_x) ** 1) * ft_u
ft_dudx2 = ((2j * np.pi * k_x) ** 2) * ft_u

It bothers me verey much because I then want to expand to partial derivatives so for loops are not an option an I must use the "power method"

If you have any idea why it doesn't work as expected, any help would be appreciated !

Versions

Collecting environment information...
PyTorch version: 2.2.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Famille
GCC version: (MinGW.org GCC-6.3.0-1) 6.3.0

Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.2 (tags/v3.12.2:6abddd9, Feb 6 2024, 21:26:36) [MSC v.1937 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-11-10.0.22631-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3050 Ti Laptop GPU
Nvidia driver version: 546.30
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9

CurrentClockSpeed=3201

DeviceID=CPU0

Family=107

L2CacheSize=4096

L2CacheSpeed=

Manufacturer=AuthenticAMD

MaxClockSpeed=3201

Name=AMD Ryzen 7 5800H with Radeon Graphics

ProcessorType=3

Revision=20480

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] optree==0.11.0
[pip3] torch==2.2.2+cu121
[pip3] torchopt==0.7.3
[pip3] torchviz==0.0.2
[conda] Could not collect