Let $\mathbf{x}$ be a column vector of $\mathbb{R}^n$. Now, if
$$ f(x) = \text{diag}(x)^2 Ax$$
what is the derivative of $f$ with respect to $x$?
For now, I looked at different formulations for the $\text{diag}(x)$ function and I came up with the following:
\begin{align} f(x) &= \text{diag}(x)\text{diag}(x)Ax\\ &= x \, \circ \, x \, \circ Ax\\ &= \text{diag}(x)\text{diag}(Ax)x = \text{diag}(x)\text{diag}(x)Ax. \end{align} Then, the derivative should be given by: \begin{equation} \frac{\partial f(x)}{\partial x} = 2\text{diag}(x)\text{diag}(Ax) + \text{diag}(x)^2 A. \end{equation}
Since I didn't succeed to find question involving $\text{diag}(x)^2$ and the closest being Derivative of a function involving diagonal matrix, I was wondering if this is correct.
Also, is the following generalisation true? \begin{equation} \frac{\partial \text{diag}(x)^n Ax}{\partial x} = n\text{diag}(x)^{n-1}\text{diag}(Ax) + \text{diag}(x)^n A \end{equation}