Jacobian of a matrix-valued function

I am fairly sure the very reason my problem arises in the first place is that I am going about it wrong, or have gotten some basic concept wrong, so any answer to point out my misunderstanding would be most welcome.

I think what I need amounts to the following: I want to get the 'Jacobian matrix' of a function $P:\mathbb{R}^n \rightarrow\mathbb{R}^{p\times q}$. As a simple example, if $U_1,\dots,U_n$ are known $\mathbb{R}^{p\times q}$ matrices, $P$ could be something such as:

$$P(\mathbf{w}) = \sum_{i=1}^nw_iU_i$$

Where $\mathbf{w} = (w_1,\dots,w_n)$. The partial derivative of $P$ with regard to any coordinate of input $\mathbf{w}$ is straightforward enough, but my problem is that I would like to have a whole Jacobian matrix to perform further computation. Except a 'three-dimensional' $\mathbb{R}^{(p\times q) \times n}$ Jacobian seems wrong, and I would have no idea how to do anything with it anyway (e.g. a simple dot product with any regular, two-dimensional matrix).

I had the intuition that the solution somehow relied on considering the Jacobian of $P$ as having $pq$ lines, but this would mean considering vectorized $(U_i)$'s. In my case, I don't see how this would work, as my $L$ function (see below) very much makes use of the $p\times q$ structure of its inputs. At the very least, I feel like it would then require to properly introduce the $vec$ transform and take it into account when computing the derivatives?

Context:

I have a function, $\mathcal{E}: \mathbb{R}^n \rightarrow\mathbb{R}$, whose gradient I want to compute. $\mathcal{E}$ itself can be written as:

$$\mathcal{E}(\mathbf{w}) = L(P(\mathbf{w}), X)$$

Where $X\in\mathbb{R}^{p\times q}$ is a known constant matrix, $L:\left(\mathbb{R}^{p\times q},\mathbb{R}^{p\times q}\right)\rightarrow\mathbb{R}$, and $P:\mathbb{R}^n \rightarrow\mathbb{R}^{p\times q}$ as above.

Now, if $P$ was a $\mathbb{R}^p$-valued function and $X\in\mathbb{R}^p$ instead, I would simply use the chain rule to get my gradient, i.e.:

$$\nabla\mathcal{E}(\mathbf{w}) = \left[\partial P(\mathbf{w})\right]^T\nabla L(P(\mathbf{w}),X)$$

In which case the (transposed) Jacobian matrix in the first term would be in $\mathbb{R}^{n\times p}$, the second term in $\mathbb{R}^p$, and my gradient in $\mathbb{R}^n$ as expected. But since $P$ is matrix-valued here, I don't know how to proceed.

Lastly, just in case this can help figuring out how I got there: in my case, I can think of $\mathbf{w}$ as a vector of weights, of the $(U_i)$ and $X$ as images and of $L$ as a loss function between my weighted transformation of $(U_i)$'s and a reference image $X$. $p$ would then be the number of pixels, and $q$ would be, say, $3$ if I am using RGB encoding - $L$ thus makes use of the $p\times q$ structure of its inputs, hence my reluctance to simply consider vectorized $(U_i)$.

There is no trick -- you've hit upon the three ways to deal with the situation: either work in coordinates, or write down the derivative of $$\mathcal{E}$$ as a rank three tensor that takes in matrices and outputs vectors, or "vectorize"/flatten your matrix before taking the derivative.
Sometimes $$L$$ has a special form that makes it easier to use one of these approaches; for example if $$L$$ is the squared Frobenius norm $$L(M,N) = \operatorname{tr}[ (M-N)^T(M-N)]$$ then it is easy to "vectorize" $$L$$; or alternatively $$dL = \operatorname {tr}\left[(dM-dN)^T(M-N) + (M-N)^T(dM-dN)\right] = 2(M-N) : (dM-dN)$$ where $$:$$ is the Frobenius product. It's not a Jacobian in matrix form (which as you note doesn't exist), but it's the best you can get.
• Thanks a bunch. Am I correct in assuming you mean the derivative of $P$ (instead of $\mathcal{E}$ using my notations from the OP) can be written as a rank three tensor though? Commented Jun 13, 2016 at 8:46