Consider the square block matrix $A\in \mathbb{R}^{mn\times mn}$ having $n^2$ blocks denoted by $A_{ij} \in \mathbb{R}^{m\times m}, i,j=1,...,n$. I need to compute the following contractions:
The block-trace $$ \sum_{i=1}^{n}A_{ii},$$ and the trace of each block $$\text{tr}(A_{ij}).$$
Suppose that I cannot access the individual elements of $A$, since I don't have the explicit representation of the matrix. The only way that I can obtain information from the matrix is by multiplying a vector $x \in\mathbb{R}^{mn}$ on the right $y=Ax$. Which vectors should I choose and how do I need to combine them in order to obtain these two contractions efficiently?
It is useful to characterize the trace of an $n \times n$ matrix as $$ \operatorname{tr}(M) = \sum_{i=1}^n e_i^TAe_i $$ where $e_1,\dots,e_n$ are the standard basis of $\Bbb R^n$ (actually, we can take them to be any orthonormal basis).
Let $\otimes$ denote the Kronecker product. Your partial traces can be expressed as $$ \sum_{i=1}^n A_{ii} = \sum_{i=1}^n (e_i \otimes I_m)^TA(e_i \otimes I_m)\\ \pmatrix{ \operatorname{tr}A_{11} & \cdots & \operatorname{tr}A_{1n}\\ \vdots & \ddots & \vdots\\ \operatorname{tr} A_{n1} & \cdots & \operatorname{tr}A_{nn} }= \sum_{j=1}^n (I_n \otimes e_j)^TA(I_n \otimes e_j) $$ Where $I_n$ denotes the $n \times n$ identity matrix.