I would like to know if I am computing a derivative correctly. Suppose I have the function $f\big(A vec(X)\big)$, where $vec(X)$ is the vectorization of matrix $X$, for some smooth function $f(\cdot)$. I want to compute the derivative wrt to $X$.
Using the chain rule:
$$ \frac{\partial f\big(A vec(X)\big)}{\partial X} = \frac{\partial f\big(A vec(X)\big)}{\partial A vec(X)} \cdot \frac{\partial Avec(X)}{\partial vec(X)} \cdot \frac{\partial vec(X)}{\partial X}$$
And this should be it, correct? For one example I want, I have that $f(\mathbf{z}) = || \mathbf{z} ||_{2}^{2}$. Therefore:
$$ \frac{\partial ||A vec(X)||_{2}^{2}}{\partial X} = 2\big(A vec(X)\big)^{T} A I$$
Am I correct?
Let's denote the reshaping of vectors/matrices by lowercase/uppercase letters and the functions $$\eqalign{ x &= {\rm vec}(X) &\Longleftrightarrow X &= {\rm Mat}(x) \cr }$$ It will also be convenient to use a colon to denote the trace/Frobenius product, i.e. $$\eqalign{A:B = {\rm tr}(A^TB)}$$ Given the vectors $$\eqalign{ z &= Ax \cr y &= A^Tz \cr }$$ and your example function, let's find the gradient wrt $X$ $$\eqalign{ f &= \|z\|^2 = z:z \cr df &= 2z:dz \cr&= 2z:A\,dx \cr&= 2A^Tz:dx \cr&= 2y:dx \cr &= 2\,{\rm Mat}(y):dX \cr &= 2\,Y:dX \cr \frac{\partial f}{\partial X} &= 2\,Y = 2\,{\rm Mat}(A^TAx) \cr }$$