tensorflow_probability/python/sts/decomposition.py
def _split_covariance_into_marginals(covariance, block_sizes):
  """Split a covariance matrix into block-diagonal marginals of given sizes."""
  start_dim = 0
  marginals = []
  for size in block_sizes:
    end_dim = start_dim + size
    marginals.append(covariance[..., start_dim:end_dim, start_dim:end_dim])
    start_dim = end_dim
  return marginals
def _split_covariance_into_marginals(covariance, block_sizes):
  """Split a covariance matrix into block-diagonal marginals of given sizes."""
  start_dim = 0
  marginals = []
  for size in block_sizes:
    end_dim = start_dim + size
    marginals.append(covariance[..., start_dim:end_dim, start_dim:end_dim])
    start_dim = end_dim
  return marginals
Split a covariance matrix into block-diagonal marginals of given sizes.