docs: def expected_diag_regression_log_prob(A, sigmasq, stats):
      """
      Expected log likelihood of p(y | x) where
  
          y_{n,d} ~ N(a_d^\trans x_n, sigma_d^2)
  
      and expectation is wrt q(y,x).  We only need expected
      sufficient statistics E[yy.T], E[yx.T], E[xx.T], and n,
      where n is the number of observations.
  
      :param A:      regression matrix
      :param sigma:  diagonal observation variance
      :param stats:  tuple (E[yy.T], E[yx.T], E[xx.T], mask)
      :return:       E[log p(y | x)]
      """
      D_out, D_in = A.shape
      assert sigmasq.shape == (D_out,)
  
      ysq, yxT, xxT, mask = stats[-4:]
      T = ysq.shape[0]
      assert ysq.shape == (T, D_out)
      assert yxT.shape == (T, D_out, D_in)
  
      # xxT has different shapes depending on whether or not data is missing
      # with missing data, it is (T, Dout, Din, Din) since we need to mask
      # off certain xxT pairs.  If there's no mask, it's always the same for
      # every output dimension.  To make it easy, we just broadcast along
      # the Dout dimension for the no-missing-data case.
      if xxT.shape == (T, D_in, D_in):
          xxT = xxT[:,None,:,:]
      else:
          assert xxT.shape == (T, D_out, D_in, D_in)
  
      assert mask.shape == (T, D_out)
  
      AAT = np.array([np.outer(a, a) for a in A])
      n = mask.sum(1)
  
      J_node = AAT[None, :, :, :] / sigmasq[None, :, None, None]
      h_node = (mask / sigmasq)[:,:,None] * A[None, :, :]
  
      out = -1 / 2. * np.sum(J_node * xxT, axis=(1, 2, 3))
      out += np.sum(h_node * yxT, axis=(1, 2))
      out += -1 / 2. * np.sum(mask / sigmasq * ysq, axis=1)
      out += -n / 2. * np.log(2 * np.pi)
      out += -1 / 2. * np.sum(mask * np.log(sigmasq), axis=1)
      assert out.shape == (T,)
      return out