Fast execution of function which resembles matrix multiplication
I think you can rearrange your equation from

\$
AB_{ija} = \frac{1}{2}(A_{ia} + B_{ja})
\$

\$
\begin{align}
C_{ij} = \sum_a &amp; A_{ia}\cdot \log\left(\frac{A_{ia}}{AB_{ija}}\right) + B_{ja}\cdot \log\left(\frac{B_{ja}}{AB_{ija}}\right) + \\
&amp;+ (1-A_{ia})\cdot \log\left(\frac{1-A_{ia}}{1-AB_{ija}}\right) + (1-B_{ja})\cdot \log\left(\frac{1-B_{ja}}{1-AB_{ija}}\right)
\end{align}
\$

to

\$
\begin{align}
C_{ij} = \sum_a &amp; A_{ia} \cdot \log A_{ia} + B_{ja}\cdot \log B_{ja} - 2AB_{ija}\cdot \log AB_{ija}\\
&amp;+ (1-A_{ia}) \cdot \log (1-A_{ia}) + (1-B_{ja})\cdot \log (1-B_{ja}) - 2(1-AB_{ija})\cdot \log (1-AB_{ija})\\
= \sum_a &amp; A_{ia}\cdot \log A_{ia} + \sum_a B_{ja}\cdot \log B_{ja} - 2\sum_a AB_{ija}\cdot \log AB_{ija} + \dots
\end{align}
\$

This will at least save you some divisions, which can be a bit slower. But the crucial part is to see that you can compute the sums separately. The sums with only \$A\$ or \$B\$ are a lot less computation, which you would otherwise repeat, and only the combined terms are (computationally) hard.

It also means that there is probably a  numpy way to eliminate at least some of your for loops, which might give you a bit of a speedup, although numba will negate some of that (by already having sped up the for loops). The hard part is then to use the numpy broadcasting correctly to achieve this, which is left as an exercise for now :).