class_docstr: The multivariate normal distribution on `R^k`.
  
  The Multivariate Normal distribution is defined over `R^k` and parameterized
  by a (batch of) length-`k` `loc` vector (aka 'mu') and a (batch of) `k x k`
  `scale` matrix; `covariance = scale @ scale.T` where `@` denotes
  matrix-multiplication.
  
  #### Mathematical Details
  
  The probability density function (pdf) is,
  
  ```none
  pdf(x; loc, scale) = exp(-0.5 ||y||**2) / Z,
  y = inv(scale) @ (x - loc),
  Z = (2 pi)**(0.5 k) |det(scale)|,
  ```
  
  where:
  
  * `loc` is a vector in `R^k`,
  * `scale` is a linear operator in `R^{k x k}`, `cov = scale @ scale.T`,
  * `Z` denotes the normalization constant, and,
  * `||y||**2` denotes the squared Euclidean norm of `y`.
  
  A (non-batch) `scale` matrix is:
  
  ```none
  scale = diag(scale_diag)
  ```
  
  where:
  
  * `scale_diag.shape = [k]`
  
  Additional leading dimensions (if any) will index batches.
  
  If `scale_diag` is `None`, then `scale` is the Identity matrix.
  
  The MultivariateNormal distribution is a member of the [location-scale
  family](https://en.wikipedia.org/wiki/Location-scale_family), i.e., it can be
  constructed as,
  
  ```none
  X ~ MultivariateNormal(loc=0, scale=1)   # Identity scale, zero shift.
  Y = scale @ X + loc
  ```
  
  #### Examples
  
  ```python
  tfd = tfp.distributions
  
  # Initialize a single 2-variate Gaussian.
  mvn = tfd.MultivariateNormalDiag(
      loc=[1., -1],
      scale_diag=[1, 2.])
  
  mvn.mean()
  # ==> [1., -1]
  
  mvn.stddev()
  # ==> [1., 2]
  
  # Evaluate this on an observation in `R^2`, returning a scalar.
  mvn.prob([-1., 0])  # shape: []
  
  # Initialize a 2-batch of 3-variate Gaussians.
  mvn = tfd.MultivariateNormalDiag(
      loc=[[1., 2, 3],
           [11, 22, 33]],           # shape: [2, 3]
      scale_diag=[[1., 2, 3],
                  [0.5, 1, 1.5]])  # shape: [2, 3]
  
  # Evaluate this on a two observations, each in `R^3`, returning a length-2
  # vector.
  x = [[-1., 0, 1],
       [-11, 0, 11.]]   # shape: [2, 3].
  mvn.prob(x)    # shape: [2]
  ```
class_docstr_tok: ["The","multivariate","normal","distribution","on","`","R^k","`",".","The","Multivariate","Normal","distribution","is","defined","over","`","R^k","`","and","parameterized","by","a","(","batch","of",")","length-","`","k","`","`","loc","`","vector","(","aka","'mu","'",")","and","a","(","batch","of",")","`","k","x","k","`","`","scale","`","matrix",";","`","covariance","=","scale","@","scale.T","`","where","`","@","`","denotes","matrix-multiplication",".","#","#","#","#","Mathematical","Details","The","probability","density","function","(","pdf",")","is",",","``","`","none","pdf","(","x",";","loc",",","scale",")","=","exp","(","-0.5","||y||","*","*","2",")","\/","Z",",","y","=","inv","(","scale",")","@","(","x","-","loc",")",",","Z","=","(","2","pi",")","*","*","(","0.5","k",")","|det","(","scale",")","|",",","``","`","where",":","*","`","loc","`","is","a","vector","in","`","R^k","`",",","*","`","scale","`","is","a","linear","operator","in","`","R^","{","k","x","k","}","`",",","`","cov","=","scale","@","scale.T","`",",","*","`","Z","`","denotes","the","normalization","constant",",","and",",","*","`","||y||","*","*","2","`","denotes","the","squared","Euclidean","norm","of","`","y","`",".","A","(","non-batch",")","`","scale","`","matrix","is",":","``","`","none","scale","=","diag","(","scale_diag",")","``","`","where",":","*","`","scale_diag.shape","=","[","k","]","`","Additional","leading","dimensions","(","if","any",")","will","index","batches",".","If","`","scale_diag","`","is","`","None","`",",","then","`","scale","`","is","the","Identity","matrix",".","The","MultivariateNormal","distribution","is","a","member","of","the","[","location-scale","family","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Location-scale_family",")",",","i.e.",",","it","can","be","constructed","as",",","``","`","none","X","~","MultivariateNormal","(","loc=0",",","scale=1",")","#","Identity","scale",",","zero","shift",".","Y","=","scale","@","X","+","loc","``","`","#","#","#","#","Examples","``","`","python","tfd","=","tfp.distributions","#","Initialize","a","single","2-variate","Gaussian",".","mvn","=","tfd.MultivariateNormalDiag","(","loc=","[","1.",",","-1","]",",","scale_diag=","[","1",",","2",".","]",")","mvn.mean","(",")","#","==",">","[","1.",",","-1","]","mvn.stddev","(",")","#","==",">","[","1.",",","2","]","#","Evaluate","this","on","an","observation","in","`","R^2","`",",","returning","a","scalar",".","mvn.prob","(","[","-1.",",","0","]",")","#","shape",":","[","]","#","Initialize","a","2-batch","of","3-variate","Gaussians",".","mvn","=","tfd.MultivariateNormalDiag","(","loc=","[","[","1.",",","2",",","3","]",",","[","11",",","22",",","33","]","]",",","#","shape",":","[","2",",","3","]","scale_diag=","[","[","1.",",","2",",","3","]",",","[","0.5",",","1",",","1.5","]","]",")","#","shape",":","[","2",",","3","]","#","Evaluate","this","on","a","two","observations",",","each","in","`","R^3","`",",","returning","a","length-2","#","vector",".","x","=","[","[","-1.",",","0",",","1","]",",","[","-11",",","0",",","11",".","]","]","#","shape",":","[","2",",","3","]",".","mvn.prob","(","x",")","#","shape",":","[","2","]","``","`"]
docstr: Construct Multivariate Normal distribution on `R^k`.
  
  The `batch_shape` is the broadcast shape between `loc` and `scale`
  arguments.
  
  The `event_shape` is given by last dimension of the matrix implied by
  `scale`. The last dimension of `loc` (if provided) must broadcast with this.
  
  Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is:
  
  ```none
  scale = diag(scale_diag)
  ```
  
  where:
  
  * `scale_diag.shape = [k]`, and,
  
  Additional leading dimensions (if any) will index batches.
  
  If `scale_diag` is `None`, then `scale` is the Identity matrix.
  
  Args:
    loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
      implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
      `b >= 0` and `k` is the event size.
    scale_diag: Non-zero, floating-point `Tensor` representing a diagonal
      matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
      and characterizes `b`-batches of `k x k` diagonal matrices added to
      `scale`. If `scale_diag` is `None` then `scale` is the Identity matrix.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    allow_nan_stats: Python `bool`, default `True`. When `True`,
      statistics (e.g., mean, mode, variance) use the value '`NaN`' to
      indicate the result is undefined. When `False`, an exception is raised
      if one or more of the statistic's batch members are undefined.
    experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
      summation to aggregate independent underlying log_prob values as well as
      when computing the log-determinant of the scale matrix. Doing so
      improves against the precision of a naive float32 sum. This can be
      noticeable in particular for large dimensions in float32. See CPU caveat
      on `tfp.math.reduce_kahan_sum`.
    name: Python `str` name prefixed to Ops created by this class.
docstr_tok: ["Construct","Multivariate","Normal","distribution","on","`","R^k","`",".","The","`","batch_shape","`","is","the","broadcast","shape","between","`","loc","`","and","`","scale","`","arguments",".","The","`","event_shape","`","is","given","by","last","dimension","of","the","matrix","implied","by","`","scale","`",".","The","last","dimension","of","`","loc","`","(","if","provided",")","must","broadcast","with","this",".","Recall","that","`","covariance","=","scale","@","scale.T","`",".","A","(","non-batch",")","`","scale","`","matrix","is",":","``","`","none","scale","=","diag","(","scale_diag",")","``","`","where",":","*","`","scale_diag.shape","=","[","k","]","`",",","and",",","Additional","leading","dimensions","(","if","any",")","will","index","batches",".","If","`","scale_diag","`","is","`","None","`",",","then","`","scale","`","is","the","Identity","matrix",".","Args",":","loc",":","Floating-point","`","Tensor","`",".","If","this","is","set","to","`","None","`",",","`","loc","`","is","implicitly","`","0","`",".","When","specified",",","may","have","shape","`","[","B1",",","...",",","Bb",",","k","]","`","where","`","b",">","=","0","`","and","`","k","`","is","the","event","size",".","scale_diag",":","Non-zero",",","floating-point","`","Tensor","`","representing","a","diagonal","matrix","added","to","`","scale","`",".","May","have","shape","`","[","B1",",","...",",","Bb",",","k","]","`",",","`","b",">","=","0","`",",","and","characterizes","`","b","`","-batches","of","`","k","x","k","`","diagonal","matrices","added","to","`","scale","`",".","If","`","scale_diag","`","is","`","None","`","then","`","scale","`","is","the","Identity","matrix",".","validate_args",":","Python","`","bool","`",",","default","`","False","`",".","When","`","True","`","distribution","parameters","are","checked","for","validity","despite","possibly","degrading","runtime","performance",".","When","`","False","`","invalid","inputs","may","silently","render","incorrect","outputs",".","allow_nan_stats",":","Python","`","bool","`",",","default","`","True","`",".","When","`","True","`",",","statistics","(","e.g.",",","mean",",","mode",",","variance",")","use","the","value","'","`","NaN","`","'","to","indicate","the","result","is","undefined",".","When","`","False","`",",","an","exception","is","raised","if","one","or","more","of","the","statistic","'s","batch","members","are","undefined",".","experimental_use_kahan_sum",":","Python","`","bool","`",".","When","`","True","`",",","we","use","Kahan","summation","to","aggregate","independent","underlying","log_prob","values","as","well","as","when","computing","the","log-determinant","of","the","scale","matrix",".","Doing","so","improves","against","the","precision","of","a","naive","float32","sum",".","This","can","be","noticeable","in","particular","for","large","dimensions","in","float32",".","See","CPU","caveat","on","`","tfp.math.reduce_kahan_sum","`",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class","."]
code: def __init__(
      self,
      loc=None,
      scale_diag=None,
      validate_args=False,
      allow_nan_stats=True,
      experimental_use_kahan_sum=False,
      name="MultivariateNormalDiag",
  ):
      """Construct Multivariate Normal distribution on `R^k`.
  
      The `batch_shape` is the broadcast shape between `loc` and `scale`
      arguments.
  
      The `event_shape` is given by last dimension of the matrix implied by
      `scale`. The last dimension of `loc` (if provided) must broadcast with this.
  
      Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is:
  
      ```none
      scale = diag(scale_diag)
      ```
  
      where:
  
      * `scale_diag.shape = [k]`, and,
  
      Additional leading dimensions (if any) will index batches.
  
      If `scale_diag` is `None`, then `scale` is the Identity matrix.
  
      Args:
        loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
          implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
          `b >= 0` and `k` is the event size.
        scale_diag: Non-zero, floating-point `Tensor` representing a diagonal
          matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
          and characterizes `b`-batches of `k x k` diagonal matrices added to
          `scale`. If `scale_diag` is `None` then `scale` is the Identity matrix.
        validate_args: Python `bool`, default `False`. When `True` distribution
          parameters are checked for validity despite possibly degrading runtime
          performance. When `False` invalid inputs may silently render incorrect
          outputs.
        allow_nan_stats: Python `bool`, default `True`. When `True`,
          statistics (e.g., mean, mode, variance) use the value '`NaN`' to
          indicate the result is undefined. When `False`, an exception is raised
          if one or more of the statistic's batch members are undefined.
        experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
          summation to aggregate independent underlying log_prob values as well as
          when computing the log-determinant of the scale matrix. Doing so
          improves against the precision of a naive float32 sum. This can be
          noticeable in particular for large dimensions in float32. See CPU caveat
          on `tfp.math.reduce_kahan_sum`.
        name: Python `str` name prefixed to Ops created by this class.
  
      """
      parameters = dict(locals())
      with tf.name_scope(name) as name:
          dtype = dtype_util.common_dtype(
              [loc, scale_diag], dtype_hint=tf.float32
          )
          loc = tensor_util.convert_nonref_to_tensor(
              loc, name="loc", dtype=dtype
          )
          scale_diag = tensor_util.convert_nonref_to_tensor(
              scale_diag, name="scale_diag", dtype=dtype
          )
          if scale_diag is not None:
              diag_cls = (
                  KahanLogDetLinOpDiag
                  if experimental_use_kahan_sum
                  else tf.linalg.LinearOperatorDiag
              )
              scale = diag_cls(
                  diag=scale_diag,
                  is_non_singular=True,
                  is_self_adjoint=True,
                  is_positive_definite=False,
              )
          else:
              # Deprecated behavior; breaks variable-safety rules by calling
              # `tf.shape(loc)`.
              num_rows = tf.compat.dimension_value(loc.shape[-1])
              if num_rows is None:
                  num_rows = tf.shape(loc)[-1]
              scale = tf.linalg.LinearOperatorIdentity(
                  num_rows=num_rows,
                  dtype=dtype,
                  is_self_adjoint=True,
                  is_positive_definite=True,
                  assert_proper_shapes=validate_args,
              )
  
          super(MultivariateNormalDiag, self).__init__(
              loc=loc,
              scale=scale,
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              experimental_use_kahan_sum=experimental_use_kahan_sum,
              name=name,
          )
          self._parameters = parameters
          self._scale_diag = scale_diag
  
file_path: probability/tensorflow_probability/python/distributions/mvn_diag.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.math import generic
  from tensorflow.python.ops.linalg import linear_operator
label_desc: Use image node_id 1 to create a new MultivariateNormalDiag object from inherited base classes: mvn_linear_operator with example: obj = MultivariateNormalDiag(loc, scale_diag, validate_args, allow_nan_stats, experimental_use_kahan_sum, name)