Hi,
If you use the same layer that has some parameters multiple times. The gradient for these weights will be the sum of the gradients for each use of it (as you would expect mathematically).
Your function look ok from the outside. What “seems incorrect”? Do you have a sample input and expected output that shows the issue?