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)$.
