The current MultivariateNormal implementation is adding a significant overhead to my code when I use large batch sizes and after looking into the source, it seems the main cause is in one helper function used in the log_prob method, which runs an explicit for loop for a batch operation:

  
      github.com
  
  
    pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/multivariate_normal.py#L55 8

return flat_bmat_inv.reshape(bmat.shape)




def _batch_trtrs_lower(bb, bA):
"""
Applies `torch.trtrs` for batches of matrices. `bb` and `bA` should have
the same batch shape.
"""
flat_b = bb.reshape((-1,) + bb.shape[-2:])
flat_A = bA.reshape((-1,) + bA.shape[-2:])
flat_X = torch.stack([torch.trtrs(b, A, upper=False)[0] for b, A in zip(flat_b, flat_A)])
return flat_X.reshape(bb.shape)




def _batch_mahalanobis(bL, bx):
r"""
Computes the squared Mahalanobis distance :math:`\mathbf{x}^\top\mathbf{M}^{-1}\mathbf{x}`
for a factored :math:`\mathbf{M} = \mathbf{L}\mathbf{L}^\top`.


Accepts batches for both bL and bx. They are not necessarily assumed to have the same batch
shape, but `bL` one should be able to broadcasted to `bx` one.



  
  
    
    
  
  


The other issue is that I just need a diagonal covariance for my current model, but the MultivariateNormal object is very general and runs some unnecessary computations that could be optimized for a diagonal covariance, like torch.trtrs. Would it make sense to have a MultivariateNormal implementation with some optimizations for strictly diagonal covariances? I’ve noticed there’s a new LowRankMultivariateNormal in the master branch that hasn’t made it into the stable release yet. I believe that implementation might be more suitable, the constructor takes a cov_diag explicitly, but it also takes a cov_factor, which might run some unnecessary computations for a strictly diagonal covariance as well:

  
      github.com
  
  
    pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/lowrank_multivariate_normal.py#L91 4

    Thanks to these formulas, we just need to compute the determinant and inverse of
    the small size "capacitance" matrix::
        capacitance = I + cov_factor.T @ inv(cov_diag) @ cov_factor
"""
arg_constraints = {"loc": constraints.real,
                   "cov_factor": constraints.real,
                   "cov_diag": constraints.positive}
support = constraints.real
has_rsample = True


def __init__(self, loc, cov_factor, cov_diag, validate_args=None):
    if loc.dim() < 1:
        raise ValueError("loc must be at least one-dimensional.")
    event_shape = loc.shape[-1:]
    if cov_factor.dim() < 2:
        raise ValueError("cov_factor must be at least two-dimensional, "
                         "with optional leading batch dimensions")
    if cov_factor.shape[-2:-1] != event_shape:
        raise ValueError("cov_factor must be a batch of matrices with shape {} x m"
                         .format(event_shape[0]))
    if cov_diag.shape[-1:] != event_shape:



  
  
    
    
  
  


What’s the recommended approach to create an efficient multivariate normal distribution with a strictly diagonal covariance matrix?