How is PyTorch taking the Mahalanobis Distance?

Alio: <p>What kind of approach is this to taking the Mahalanobis Distance?</p>
<p>I’m looking in the source code of PyTorch in distributions and the <a href="https://github.com/pytorch/pytorch/blob/master/torch/distributions/multivariate_normal.py" rel="noopener nofollow ugc">Mahalanobis Distance</a> does not look anything like the typical equation you see returned when looking up the equation online.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/e/5e8fe16492e1f4cd224110b14fb93b5dab8867e0.png" data-download-href="https://discuss.pytorch.org/uploads/default/5e8fe16492e1f4cd224110b14fb93b5dab8867e0" title="56065046504650"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/e/5e8fe16492e1f4cd224110b14fb93b5dab8867e0_2_689x148.png" alt="56065046504650" data-base62-sha1="dux8Y0YNqBYyAdNhxBS89o4afmw" width="689" height="148" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/e/5e8fe16492e1f4cd224110b14fb93b5dab8867e0_2_689x148.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/e/5e8fe16492e1f4cd224110b14fb93b5dab8867e0_2_1033x222.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/e/5e8fe16492e1f4cd224110b14fb93b5dab8867e0.png 2x" data-dominant-color="F3F3F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">56065046504650</span><span class="informations">1346×289 31.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Can I have some insight into what is happening here?</p>
<pre><code class="lang-auto">
def _batch_mahalanobis(bL, bx):

    n = bx.size(-1)
    bx_batch_shape = bx.shape[:-1]

    # Assume that bL.shape = (i, 1, n, n), bx.shape = (..., i, j, n),
    # we are going to make bx have shape (..., 1, j,  i, 1, n) to apply batched tri.solve
    bx_batch_dims = len(bx_batch_shape)
    bL_batch_dims = bL.dim() - 2

    outer_batch_dims = bx_batch_dims - bL_batch_dims
    old_batch_dims = outer_batch_dims + bL_batch_dims
    new_batch_dims = outer_batch_dims + 2 * bL_batch_dims
    # Reshape bx with the shape (..., 1, i, j, 1, n)
    bx_new_shape = bx.shape[:outer_batch_dims]
    for (sL, sx) in zip(bL.shape[:-2], bx.shape[outer_batch_dims:-1]):
        bx_new_shape += (sx // sL, sL)
    bx_new_shape += (n,)
    bx = bx.reshape(bx_new_shape)
    # Permute bx to make it have shape (..., 1, j, i, 1, n)
    permute_dims = (list(range(outer_batch_dims)) +
                    list(range(outer_batch_dims, new_batch_dims, 2)) +
                    list(range(outer_batch_dims + 1, new_batch_dims, 2)) +
                    [new_batch_dims])
    bx = bx.permute(permute_dims)

    flat_L = bL.reshape(-1, n, n)  # shape = b x n x n
    flat_x = bx.reshape(-1, flat_L.size(0), n)  # shape = c x b x n
    flat_x_swap = flat_x.permute(1, 2, 0)  # shape = b x n x c
    M_swap = torch.linalg.solve_triangular(flat_L, flat_x_swap, upper=False).pow(2).sum(-2)  # shape = b x c
    M = M_swap.t()  # shape = c x b

    permuted_M = M.reshape(bx.shape[:-1])  # shape = (..., 1, j, i, 1)
    permute_inv_dims = list(range(outer_batch_dims))
    for i in range(bL_batch_dims):
        permute_inv_dims += [outer_batch_dims + i, old_batch_dims + i]
    reshaped_M = permuted_M.permute(permute_inv_dims)  # shape = (..., 1, i, j, 1)
    return reshaped_M.reshape(bx_batch_shape)
</code></pre>

p-adema: <p>Since I had to figure it out, here’s the unbatched version for any future readers (if you squint my handwriting away):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/9/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df.png" data-download-href="https://discuss.pytorch.org/uploads/default/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df_2_641x500.png" alt="image" data-base62-sha1="qqKI2DKZ7p6DTRPE3rrBp9yQI4n" width="641" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df_2_641x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df_2_961x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b93eb87b40571c9e72c8db1f5ae9be3d3e3a09df_2_1282x1000.png 2x" data-dominant-color="F2F0E2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1586×1237 123 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Most of the complexity in that function is wrangling L and x to work properly when batched.</p>