I am looking for the derivative of $\frac{dY}{dX}$ where $Y$ and $X$ are both a vector of $N\times 1$. And $Y$ is $$(I- A \operatorname{diag}(X))^{-1} K\ ,$$ where $K$ is a vector of $N\times 1$, and $I$ is the identity matrix of $N\times N$. A is a $N\times N$ matrix.
I get this: $$\frac{dY}{dX} = (I- A \operatorname{diag}(X))^{-1} A (I- A \operatorname{diag}(X))^{-1} K\ .$$
However, RHS dimension is $N \times 1$ and $\frac{dY}{dX}$ should be $N\times N$.
I appreciate it if you tell me where I got wrong.
$ \def\bbR#1{{\mathbb R}^{#1}} \def\e{\varepsilon} \def\E{{\cal E}} \def\n{\nabla} \def\o{{\tt1}} \def\p{\partial} \def\bx{\boxtimes} \def\LR#1{\left(#1\right)} \def\vvec#1{\operatorname{vec}\LR{#1}} \def\diag#1{\operatorname{diag}\LR{#1}} \def\Diag#1{\operatorname{Diag}\LR{#1}} \def\trace#1{\operatorname{Tr}\LR{#1}} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\hess#1#2#3{\frac{\p^2 #1}{\p #2\,\p #3^T}} \def\c#1{\color{red}{#1}} \def\m#1{\left[\begin{array}{r}#1\end{array}\right]} \def\mc#1{\left[\begin{array}{r|r}#1\end{array}\right]} \def\qiq{\quad\implies\quad} $The $\sf Khatri$-$\sf Rao$ product $(\bx)$ is defined as the columnwise Kronecker product $(\otimes)$ $$\eqalign{ F &= \LR{G\bx H} \;\iff\; f_k = \LR{g_k\otimes h_k} \\ H &= \mc{h_1&h_2&\ldots&h_n} \;\in\bbR{m\times n} \\ G &= \mc{g_1&\,g_2&\ldots&\,g_n} \;\in\bbR{q\times n} \\ F &= \mc{f_1\,&f_2&\ldots&\,f_n} \;\in\bbR{(mq)\times n} \\ }$$ or in terms of all-ones vectors $\LR{\o_n\in\bbR n}$ and the Hadamard product $(\odot)$ $$\eqalign{ G\bx H &= (G\otimes\o_m)\odot(\o_q\otimes H) \;\in \bbR{(mq)\times n} \\ }$$ It is useful in vectorizing expressions involving a diagonal matrix $$\eqalign{ &\vvec{A\:\Diag{x}\;B} = \LR{B^T\bx A} x \\ }$$
For your problem, define the following variables $$\eqalign{ X &= \Diag x &\qiq dX = \Diag{dx} \\ B &= \LR{I-AX}^{-1} &\qiq dB = BA\:dX\:B \\ }$$ then calculate the gradient of the function like so $$\eqalign{ y &= Bk \\ dy &= dB\:k \\ &= \LR{BA\:dX\:B} k \\ &= BA\:dX\;y \\ &= \vvec{BA\;dX\;y} \\ &= \LR{y^T\!\bx BA} dx \\ \grad yx &= y^T\!\bx BA \\ }$$