class_docstr: Mixture of Normal distributions with arbitrary means and arbitrary
  diagonal covariance matrices.
  
  That is, this distribution is a mixture with K components, where each
  component distribution is a D-dimensional Normal distribution with a
  D-dimensional mean parameter and a D-dimensional diagonal covariance
  matrix. The K different component means are gathered into the K x D
  dimensional parameter `locs` and the K different scale parameters are
  gathered into the K x D dimensional parameter `coord_scale`. The mixture
  weights are controlled by a K-dimensional vector of softmax logits,
  `component_logits`. This distribution implements pathwise derivatives
  for samples from the distribution.
  
  See reference [1] for details on the implementations of the pathwise
  derivative. Please consider citing this reference if you use the pathwise
  derivative in your research. Note that this distribution does not support
  dimension D = 1.
  
  [1] Pathwise Derivatives for Multivariate Distributions, Martin Jankowiak &
  Theofanis Karaletsos. arXiv:1806.01856
  
  :param torch.Tensor locs: K x D mean matrix
  :param torch.Tensor coord_scale: K x D scale matrix
  :param torch.Tensor component_logits: K-dimensional vector of softmax logits
class_docstr_tok: ["Mixture","of","Normal","distributions","with","arbitrary","means","and","arbitrary","diagonal","covariance","matrices",".","That","is",",","this","distribution","is","a","mixture","with","K","components",",","where","each","component","distribution","is","a","D-dimensional","Normal","distribution","with","a","D-dimensional","mean","parameter","and","a","D-dimensional","diagonal","covariance","matrix",".","The","K","different","component","means","are","gathered","into","the","K","x","D","dimensional","parameter","`","locs","`","and","the","K","different","scale","parameters","are","gathered","into","the","K","x","D","dimensional","parameter","`","coord_scale","`",".","The","mixture","weights","are","controlled","by","a","K-dimensional","vector","of","softmax","logits",",","`","component_logits","`",".","This","distribution","implements","pathwise","derivatives","for","samples","from","the","distribution",".","See","reference","[","1","]","for","details","on","the","implementations","of","the","pathwise","derivative",".","Please","consider","citing","this","reference","if","you","use","the","pathwise","derivative","in","your","research",".","Note","that","this","distribution","does","not","support","dimension","D","=","1",".","[","1","]","Pathwise","Derivatives","for","Multivariate","Distributions",",","Martin","Jankowiak","&","Theofanis","Karaletsos",".","arXiv:1806.01856",":","param","torch.Tensor","locs",":","K","x","D","mean","matrix",":","param","torch.Tensor","coord_scale",":","K","x","D","scale","matrix",":","param","torch.Tensor","component_logits",":","K-dimensional","vector","of","softmax","logits"]
code: def __init__(self, locs, coord_scale, component_logits):
      self.batch_mode = locs.dim() > 2
      assert coord_scale.shape == locs.shape
      assert self.batch_mode or locs.dim() == 2, (
          "The locs parameter in MixtureOfDiagNormals should be K x D dimensional "
          "(or ... x B x K x D if doing batches)"
      )
      if not self.batch_mode:
          assert (
              coord_scale.dim() == 2
          ), "The coord_scale parameter in MixtureOfDiagNormals should be K x D dimensional"
          assert (
              component_logits.dim() == 1
          ), "The component_logits parameter in MixtureOfDiagNormals should be K dimensional"
          assert component_logits.size(-1) == locs.size(-2)
          batch_shape = ()
      else:
          assert (
              coord_scale.dim() > 2
          ), "The coord_scale parameter in MixtureOfDiagNormals should be ... x B x K x D dimensional"
          assert (
              component_logits.dim() > 1
          ), "The component_logits parameter in MixtureOfDiagNormals should be ... x B x K dimensional"
          assert component_logits.size(-1) == locs.size(-2)
          batch_shape = tuple(locs.shape[:-2])
  
      self.locs = locs
      self.coord_scale = coord_scale
      self.component_logits = component_logits
      self.dim = locs.size(-1)
      self.categorical = Categorical(logits=component_logits)
      self.probs = self.categorical.probs
      super().__init__(
          batch_shape=torch.Size(batch_shape),
          event_shape=torch.Size((self.dim,)),
      )
  
file_path: pyro/pyro/distributions/diag_normal_mixture.py
imports: import math
  import torch
  from torch.autograd import Function
  from torch.autograd.function import once_differentiable
  from torch.distributions import Categorical, constraints
  from pyro.distributions.torch_distribution import TorchDistribution
  from pyro.distributions.util import sum_leftmost
label_desc: Use image node_id 1 to create a new MixtureOfDiagNormals object from inherited base classes: TorchDistribution with example: obj = MixtureOfDiagNormals(locs, coord_scale, component_logits)