General Question: Gradient of Sum

[praisethemoon]

Hi!

I came a cross this library a couple of weeks and I am toying with it for I am writing an automatic differentiation tools my self. I am comparing my results with autograd, but I am stuck with gradient of numpy ndarray sum across a given axis, here is the scenario:

import autograd.numpy as np
from autograd import elementwise_grad as grad

y = np.linspace(0, 3, 9*3).reshape(3, 3, 3)

def f(x):
    return np.sum(x, axis=0)*np.linspace(0, 3, 9).reshape(3, 3)

g = grad(f)

>>> g(y)
array([[[0.   , 0.375, 0.75 ],
        [1.125, 1.5  , 1.875],
        [2.25 , 2.625, 3.   ]],

       [[0.   , 0.375, 0.75 ],
        [1.125, 1.5  , 1.875],
        [2.25 , 2.625, 3.   ]],

       [[0.   , 0.375, 0.75 ],
        [1.125, 1.5  , 1.875],
        [2.25 , 2.625, 3.   ]]])

My question is, altering the sum axis changes obviously changes the gradient, but I did not find any formula or the corresponding code for some inspirations in order to implement it.

If you have any tips please let me know!

Cheers!

[ZoeDuan2021]

need to import np from autograd