You have calculated the gradient of the function $f(A)$
$$G=\frac{\partial f}{\partial A}$$
with no constraints on the matrix variable, and now you wish to constrain $A$ to be diagonal, i.e.
$$A={\rm Diag}(a)$$
Start with a differential in terms of $dA$, then change the variable to $da$
$$\eqalign{
df &= G:dA \cr
&= G:{\rm Diag}(da) \cr
&= {\rm diag}(G):da \cr
\frac{\partial f}{\partial a}
&= {\rm diag}(G) \cr
&= {\rm diag}\Big(\frac{\partial f}{\partial A}\Big) \cr
}$$
where
$\,\,\,:\,\,$ is a product notation for the trace $\,\,A:B={\rm Tr}(A^TB)$
$\,\,\,{\rm Diag}()$ generates a diagonal matrix from the input vector
${\,\,\,\rm diag}()$ extracts the diagonal of a matrix into an output vector