class_docstr: Funsor representing a batched Gaussian log-density function.
  
  Gaussians are the internal representation for joint and conditional
  multivariate normal distributions and multivariate normal likelihoods.
  Mathematically, a Gaussian represents the quadratic log density function::
  
      f(x) = -0.5 * || x @ prec_sqrt - white_vec ||^2
           = -0.5 * < x @ prec_sqrt - white_vec | x @ prec_sqrt - white_vec >
           = -0.5 * < x | prec_sqrt @ prec_sqrt.T | x>
             + < x | prec_sqrt | white_vec > - 0.5 ||white_vec||^2
  
  Internally Gaussians use a square root information filter (SRIF)
  representation consisting of a square root of the precision matrix
  ``prec_sqrt`` and a vector in the whitened space ``white_vec``. This
  representation allows space-efficient construction of Gaussians with
  incomplete information, i.e. with zero eigenvalues in the precision matrix.
  These incomplete log densities arise when making low-dimensional
  observations of higher-dimensional hidden state. Sampling and
  marginalization are supported only for full-rank Gaussians or full-rank
  subsets of Gaussians. See the :meth:`rank` and :meth:`is_full_rank`
  properties.
  
  .. note:: :class:`Gaussian` s are not normalized probability distributions,
      rather they are canonicalized to evaluate to zero log density at their
      maximum: ``f(prec_sqrt \ white_vec) = 0``. Not only are Gaussians
      non-normalized, but they may be rank deficient and non-normalizable, in
      which case sampling and marginalization are supported only un full-rank
      subsets of variables.
  
  :param torch.Tensor white_vec: An batched white noise vector, where
      ``white_vec = prec_sqrt.T @ mean``. Alternatively you can specify one
      of the kwargs ``mean`` or ``info_vec``, which will be converted to
      ``white_vec``.
  :param torch.Tensor prec_sqrt: A batched square root of the positive
      semidefinite precision matrix. This need not be square, and typically
      has shape ``prec_sqrt.shape == white_vec.shape[:-1] + (dim, rank)``,
      where ``dim`` is the total flattened size of real inputs and
      ``rank = white_vec.shape[-1]``.  Alternatively you can specify one of
      the kwargs ``precision``, ``covariance``, or ``scale_tril``, which will
      be converted to ``prec_sqrt``.
  :param OrderedDict inputs: Mapping from name to
      :class:`~funsor.domains.Domain` .
class_docstr_tok: ["Funsor","representing","a","batched","Gaussian","log-density","function",".","Gaussians","are","the","internal","representation","for","joint","and","conditional","multivariate","normal","distributions","and","multivariate","normal","likelihoods",".","Mathematically",",","a","Gaussian","represents","the","quadratic","log","density","function",":",":","f","(","x",")","=","-0.5","*","||","x","@","prec_sqrt","-","white_vec","||^2","=","-0.5","*","<","x","@","prec_sqrt","-","white_vec","|","x","@","prec_sqrt","-","white_vec",">","=","-0.5","*","<","x","|","prec_sqrt","@","prec_sqrt.T","|","x",">","+","<","x","|","prec_sqrt","|","white_vec",">","-","0.5","||white_vec||^2","Internally","Gaussians","use","a","square","root","information","filter","(","SRIF",")","representation","consisting","of","a","square","root","of","the","precision","matrix","``","prec_sqrt","``","and","a","vector","in","the","whitened","space","``","white_vec","``",".","This","representation","allows","space-efficient","construction","of","Gaussians","with","incomplete","information",",","i.e",".","with","zero","eigenvalues","in","the","precision","matrix",".","These","incomplete","log","densities","arise","when","making","low-dimensional","observations","of","higher-dimensional","hidden","state",".","Sampling","and","marginalization","are","supported","only","for","full-rank","Gaussians","or","full-rank","subsets","of","Gaussians",".","See","the",":","meth",":","`","rank","`","and",":","meth",":","`","is_full_rank","`","properties",".","..","note",":",":",":","class",":","`","Gaussian","`","s","are","not","normalized","probability","distributions",",","rather","they","are","canonicalized","to","evaluate","to","zero","log","density","at","their","maximum",":","``","f","(","prec_sqrt","\\","white_vec",")","=","0","``",".","Not","only","are","Gaussians","non-normalized",",","but","they","may","be","rank","deficient","and","non-normalizable",",","in","which","case","sampling","and","marginalization","are","supported","only","un","full-rank","subsets","of","variables",".",":","param","torch.Tensor","white_vec",":","An","batched","white","noise","vector",",","where","``","white_vec","=","prec_sqrt.T","@","mean","``",".","Alternatively","you","can","specify","one","of","the","kwargs","``","mean","``","or","``","info_vec","``",",","which","will","be","converted","to","``","white_vec","``",".",":","param","torch.Tensor","prec_sqrt",":","A","batched","square","root","of","the","positive","semidefinite","precision","matrix",".","This","need","not","be","square",",","and","typically","has","shape","``","prec_sqrt.shape","==","white_vec.shape","[",":","-1","]","+","(","dim",",","rank",")","``",",","where","``","dim","``","is","the","total","flattened","size","of","real","inputs","and","``","rank","=","white_vec.shape","[","-1","]","``",".","Alternatively","you","can","specify","one","of","the","kwargs","``","precision","``",",","``","covariance","``",",","or","``","scale_tril","``",",","which","will","be","converted","to","``","prec_sqrt","``",".",":","param","OrderedDict","inputs",":","Mapping","from","name","to",":","class",":","`","~funsor.domains.Domain","`","."]
code: def eager_reduce(self, op, reduced_vars):
      assert reduced_vars.issubset(self.inputs)
      if op is ops.logaddexp:
          # Marginalize out real variables, but keep mixtures lazy.
          assert all(v in self.inputs for v in reduced_vars)
          real_vars = frozenset(
              k for k, d in self.inputs.items() if d.dtype == "real"
          )
          reduced_reals = reduced_vars & real_vars
          reduced_ints = reduced_vars - real_vars
          if not reduced_reals:
              return None  # defer to default implementation
          if reduced_reals == real_vars:
              return self.log_normalizer.reduce(
                  ops.logaddexp, reduced_ints
              )
  
          inputs = OrderedDict(
              (k, d)
              for k, d in self.inputs.items()
              if k not in reduced_reals
          )
          int_inputs = OrderedDict(
              (k, v) for k, v in inputs.items() if v.dtype != "real"
          )
  
          # Let x = [xa xb] where xb will be marginalized out and will xa
          # remain. Following the formula in _compress_rank, we can rewrite
          # the joint Gaussian as a Gaussian in xb plus a term C that does
          # not depend on xb:
          #
          #   log integral exp G([xa xb]; w, [ Pa ]) dxb
          #                                  [ Pb ]
          #     =: G(xb; wb, Qb).log_normalizer + C
          #
          # In normalizable models, rank >= dim(xb), so we can choose Qb to
          # be the Cholesky square root, making it easy to compute a
          # determinant and solve for wb.
          #
          #    Qb = chol(Pb Pb')
          #   wb' = Qb \ Pb (w - xa Pa)'
          #
          # Next we match moments of C to a Gaussian in xa:
          #
          #   C = 1/2 (wb wb' - w w')  # from _compress_rank
          #     = 1/2 (xa Pa - w) Pb' inv(Qb Qb') Pb (xa Pa - w)'
          #     - 1/2 (xa Pa - w) (xa Pa - w)'
          #     =: G(xa; wa, Qa)
          #
          # whence  Qa = Pa S  and  wa = w S, where S is a square root of the
          # rank-by-rank projection matrix (S = S S' by idempotence):
          #
          #   S S' = I - Pb' inv(Pb Pb') Pb = I - (Qb\Pb)' (Qb\Pb) = S
          #
          # Note if rank == dim(xb), then the projection matrix is zero,
          # and the Gaussian G(xa; wa, Qa) is zero can be dropped.
          b, a = _split_real_inputs(
              self.inputs, reduced_vars, self.white_vec
          )
          prec_sqrt_a = self.prec_sqrt[..., a, :]
          prec_sqrt_b = self.prec_sqrt[..., b, :]
          dim_b = prec_sqrt_b.shape[-2]
          if self.rank < dim_b:
              raise ValueError(
                  f"Too little information to marginalize over {set(reduced_vars)}. "
                  "Consider adding a prior."
              )
          precision_chol_b = ops.cholesky(
              _mmt(prec_sqrt_b)
          )  # assume full rank
          result = self._marginalize_after_split(
              inputs,
              int_inputs,
              prec_sqrt_b,
              prec_sqrt_a,
              precision_chol_b,
          )
          return result.reduce(ops.logaddexp, reduced_ints)
  
      elif op is ops.add:
          # Fuse Gaussians along a plate. Compare to eager_add_gaussian_gaussian().
          inputs = OrderedDict()
          old_ints = OrderedDict()
          new_ints = OrderedDict()
          kept_perm = []
          reduced_perm = []
          for i, (k, v) in enumerate(self.inputs.items()):
              if k not in reduced_vars:
                  inputs[k] = v
              if v.dtype == "real":
                  if v in reduced_vars:
                      raise ValueError(
                          f"Cannot sum along a real dimension: {repr(v)}"
                      )
              else:
                  old_ints[k] = v
                  if k in reduced_vars:
                      reduced_perm.append(i)
                  else:
                      kept_perm.append(i)
                      new_ints[k] = v
          n = len(kept_perm) + len(reduced_perm)
  
          # The square root information filter fuses via transpose and reshape.
          perm = kept_perm + reduced_perm + [n]
          white_vec = ops.permute(self.white_vec, perm)
          white_vec = white_vec.reshape(
              white_vec.shape[: len(kept_perm)] + (-1,)
          )
          perm = kept_perm + [n] + reduced_perm + [n + 1]
          prec_sqrt = ops.permute(self.prec_sqrt, perm)
          prec_sqrt = prec_sqrt.reshape(
              prec_sqrt.shape[: len(kept_perm) + 1] + (-1,)
          )
          assert prec_sqrt.shape[:-2] == white_vec.shape[:-1]
          assert prec_sqrt.shape[-1] == white_vec.shape[-1]
  
          return Gaussian(white_vec, prec_sqrt, inputs)
  
      return None
  
imports: import math
  from collections import OrderedDict, defaultdict
  from contextlib import contextmanager
  from functools import reduce
  import funsor.ops
  from funsor.affine import affine_inputs, extract_affine, is_affine
  from funsor.delta import Delta
  from funsor.domains import Real, Reals
  from funsor.interpretations import compress_gaussians
  from funsor.ops import AddOp, SubOp
  from funsor.tensor import Tensor, align_tensor, align_tensors
  from funsor.terms import Binary, Funsor, FunsorMeta, Number, Slice, Subs, Variable, eager, reflect
  from funsor.util import broadcast_shape, get_tracing_state, lazy_property
label_desc: Use image node_id 20 for calling the Gaussian obj's underlying member method code with example usage: obj.eager_reduce(op, reduced_vars) and returns: None, result, None, self, Gaussian