def multigaussian_loss(preds, targets, ngauss=1):  # pylint: disable=invalid-name
  """Returns a mixture of gaussians loss.

  Args:
    preds: <tbd>
    targets: <tbd>
    ngauss: <tbd>
  """
  ndims = targets.shape[-1]
  logits = preds[:, :ngauss]
  mus = preds[:, ngauss:ngauss*(ndims + 1)]
  sigmas = preds[:, ngauss(ndims + 1):]
  sigmas = sigmas * sigmas + 1e-6  # Make positive.
  loglogits = logits - fastmath.logsumexp(logits, axis=-1, keepdims=True)
  mus = jnp.reshape(mus, [-1, ngauss, ndims])
  sigmas = jnp.reshape(sigmas, [-1, ngauss, ndims])
  targets = jnp.reshape(targets, [-1, 1, ndims])
  glogprobs = log_gaussian_diag_pdf(targets, mus, sigmas)
  return fastmath.logsumexp(loglogits + glogprobs, axis=-1)
Returns a mixture of gaussians loss.

  Args:
    preds: <tbd>
    targets: <tbd>
    ngauss: <tbd>
  
trax/layers/core.py