class_docstr: Posterior predictive in a conjugate Multi-task GP regression model.
docstr: Construct a MultiTaskGaussianProcessRegressionModelWithCholesky instance.
  
  Args:
    kernel: `MultiTaskKernel`-like instance representing the GP's covariance
      function.
    observation_index_points: (Nested) `float` `Tensor` representing finite
      collection, or batch of collections, of points in the index set for
      which some data has been observed. Shape (of each nested component) has
      the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of
      feature dimensions and must equal `kernel.feature_ndims` (or its
      corresponding nested component), and `e` is the number (size) of index
      points in each batch. `[b1, ..., bB, e]` must be broadcastable with the
      shape of `observations`, and `[b1, ..., bB]` must be broadcastable with
      the shapes of all other batched parameters (`kernel.batch_shape`,
      `index_points`, etc).
    observations: `float` `Tensor` representing collection, or batch of
      collections, of observations corresponding to
      `observation_index_points`. Shape has the form `[b1, ..., bB, e, t]`,
      which must be broadcastable with the batch and example shapes of
      `observation_index_points`. The batch shape `[b1, ..., bB]` must be
      broadcastable with the shapes of all other batched parameters
      (`kernel.batch_shape`, `index_points`, etc.).
    observations_is_missing:  `bool` `Tensor` of shape `[..., e, t]`,
      representing a batch of boolean masks.  When
      `observations_is_missing` is not `None`, this distribution is
      conditioned only on the observations for which the
      corresponding elements of `observations_is_missing` are `False`.
    index_points: (Nested) `float` `Tensor` representing finite collection, or
      batch of collections, of points in the index set over which the GP is
      defined.  Shape (or shape of each nested component) has the form
      `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
      dimensions and must equal `kernel.feature_ndims` (or its corresponding
      nested component) and `e` is the number (size) of index points in each
      batch. Ultimately this distribution corresponds to an `e`-dimensional
      multivariate normal. The batch shape must be broadcastable with
      `kernel.batch_shape`.
    mean_fn: Python `callable` that acts on `index_points` to produce a (batch
      of) collection of mean values at `index_points`. Takes a (nested)
      `Tensor` of shape `[b1, ..., bB, e, f1, ..., fF]` and returns a `Tensor`
      whose shape is broadcastable with `[b1, ..., bB, e, t]`, where `t` is
      the number of tasks.
    observation_noise_variance: `float` `Tensor` representing the variance of
      the noise in the Normal likelihood distribution of the model. May be
      batched, in which case the batch shape must be broadcastable with the
      shapes of all other batched parameters (`kernel.batch_shape`,
      `index_points`, etc.).
      Default value: `None`
    predictive_noise_variance: `float` `Tensor` representing the variance in
      the posterior predictive model. If `None`, we simply re-use
      `observation_noise_variance` for the posterior predictive noise. If set
      explicitly, however, we use this value. This allows us, for example, to
      omit predictive noise variance (by setting this to zero) to obtain
      noiseless posterior predictions of function values, conditioned on noisy
      observations.
    cholesky_fn: Callable which takes a single (batch) matrix argument and
      returns a Cholesky-like lower triangular factor.  Default value: `None`,
        in which case `make_cholesky_with_jitter_fn(1e-6)` is used.
    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.
      Default value: `False`.
    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.
      Default value: `False`.
    name: Python `str` name prefixed to Ops created by this class.
      Default value: 'MultiTaskGaussianProcessRegressionModel'.
    _flattened_conditional_mean_fn: Internal parameter -- do not use.
    _observation_scale: Internal parameter -- do not use.
docstr_tok: ["Construct","a","MultiTaskGaussianProcessRegressionModelWithCholesky","instance",".","Args",":","kernel",":","`","MultiTaskKernel","`","-like","instance","representing","the","GP","'s","covariance","function",".","observation_index_points",":","(","Nested",")","`","float","`","`","Tensor","`","representing","finite","collection",",","or","batch","of","collections",",","of","points","in","the","index","set","for","which","some","data","has","been","observed",".","Shape","(","of","each","nested","component",")","has","the","form","`","[","b1",",","...",",","bB",",","e",",","f1",",","...",",","fF","]","`","where","`","F","`","is","the","number","of","feature","dimensions","and","must","equal","`","kernel.feature_ndims","`","(","or","its","corresponding","nested","component",")",",","and","`","e","`","is","the","number","(","size",")","of","index","points","in","each","batch",".","`","[","b1",",","...",",","bB",",","e","]","`","must","be","broadcastable","with","the","shape","of","`","observations","`",",","and","`","[","b1",",","...",",","bB","]","`","must","be","broadcastable","with","the","shapes","of","all","other","batched","parameters","(","`","kernel.batch_shape","`",",","`","index_points","`",",","etc",")",".","observations",":","`","float","`","`","Tensor","`","representing","collection",",","or","batch","of","collections",",","of","observations","corresponding","to","`","observation_index_points","`",".","Shape","has","the","form","`","[","b1",",","...",",","bB",",","e",",","t","]","`",",","which","must","be","broadcastable","with","the","batch","and","example","shapes","of","`","observation_index_points","`",".","The","batch","shape","`","[","b1",",","...",",","bB","]","`","must","be","broadcastable","with","the","shapes","of","all","other","batched","parameters","(","`","kernel.batch_shape","`",",","`","index_points","`",",","etc.",")",".","observations_is_missing",":","`","bool","`","`","Tensor","`","of","shape","`","[","...",",","e",",","t","]","`",",","representing","a","batch","of","boolean","masks",".","When","`","observations_is_missing","`","is","not","`","None","`",",","this","distribution","is","conditioned","only","on","the","observations","for","which","the","corresponding","elements","of","`","observations_is_missing","`","are","`","False","`",".","index_points",":","(","Nested",")","`","float","`","`","Tensor","`","representing","finite","collection",",","or","batch","of","collections",",","of","points","in","the","index","set","over","which","the","GP","is","defined",".","Shape","(","or","shape","of","each","nested","component",")","has","the","form","`","[","b1",",","...",",","bB",",","e",",","f1",",","...",",","fF","]","`","where","`","F","`","is","the","number","of","feature","dimensions","and","must","equal","`","kernel.feature_ndims","`","(","or","its","corresponding","nested","component",")","and","`","e","`","is","the","number","(","size",")","of","index","points","in","each","batch",".","Ultimately","this","distribution","corresponds","to","an","`","e","`","-dimensional","multivariate","normal",".","The","batch","shape","must","be","broadcastable","with","`","kernel.batch_shape","`",".","mean_fn",":","Python","`","callable","`","that","acts","on","`","index_points","`","to","produce","a","(","batch","of",")","collection","of","mean","values","at","`","index_points","`",".","Takes","a","(","nested",")","`","Tensor","`","of","shape","`","[","b1",",","...",",","bB",",","e",",","f1",",","...",",","fF","]","`","and","returns","a","`","Tensor","`","whose","shape","is","broadcastable","with","`","[","b1",",","...",",","bB",",","e",",","t","]","`",",","where","`","t","`","is","the","number","of","tasks",".","observation_noise_variance",":","`","float","`","`","Tensor","`","representing","the","variance","of","the","noise","in","the","Normal","likelihood","distribution","of","the","model",".","May","be","batched",",","in","which","case","the","batch","shape","must","be","broadcastable","with","the","shapes","of","all","other","batched","parameters","(","`","kernel.batch_shape","`",",","`","index_points","`",",","etc.",")",".","Default","value",":","`","None","`","predictive_noise_variance",":","`","float","`","`","Tensor","`","representing","the","variance","in","the","posterior","predictive","model",".","If","`","None","`",",","we","simply","re-use","`","observation_noise_variance","`","for","the","posterior","predictive","noise",".","If","set","explicitly",",","however",",","we","use","this","value",".","This","allows","us",",","for","example",",","to","omit","predictive","noise","variance","(","by","setting","this","to","zero",")","to","obtain","noiseless","posterior","predictions","of","function","values",",","conditioned","on","noisy","observations",".","cholesky_fn",":","Callable","which","takes","a","single","(","batch",")","matrix","argument","and","returns","a","Cholesky-like","lower","triangular","factor",".","Default","value",":","`","None","`",",","in","which","case","`","make_cholesky_with_jitter_fn","(","1e-6",")","`","is","used",".","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",".","Default","value",":","`","False","`",".","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",".","Default","value",":","`","False","`",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","Default","value",":","'MultiTaskGaussianProcessRegressionModel","'",".","_flattened_conditional_mean_fn",":","Internal","parameter","--","do","not","use",".","_observation_scale",":","Internal","parameter","--","do","not","use","."]
code: def __init__(
      self,
      kernel,
      observation_index_points,
      observations,
      observations_is_missing=None,
      index_points=None,
      mean_fn=None,
      observation_noise_variance=None,
      predictive_noise_variance=None,
      cholesky_fn=None,
      validate_args=False,
      allow_nan_stats=False,
      name="MultiTaskGaussianProcessRegressionModelWithCholesky",
      _flattened_conditional_mean_fn=None,
      _observation_scale=None,
  ):
      """Construct a MultiTaskGaussianProcessRegressionModelWithCholesky instance.
  
      Args:
        kernel: `MultiTaskKernel`-like instance representing the GP's covariance
          function.
        observation_index_points: (Nested) `float` `Tensor` representing finite
          collection, or batch of collections, of points in the index set for
          which some data has been observed. Shape (of each nested component) has
          the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of
          feature dimensions and must equal `kernel.feature_ndims` (or its
          corresponding nested component), and `e` is the number (size) of index
          points in each batch. `[b1, ..., bB, e]` must be broadcastable with the
          shape of `observations`, and `[b1, ..., bB]` must be broadcastable with
          the shapes of all other batched parameters (`kernel.batch_shape`,
          `index_points`, etc).
        observations: `float` `Tensor` representing collection, or batch of
          collections, of observations corresponding to
          `observation_index_points`. Shape has the form `[b1, ..., bB, e, t]`,
          which must be broadcastable with the batch and example shapes of
          `observation_index_points`. The batch shape `[b1, ..., bB]` must be
          broadcastable with the shapes of all other batched parameters
          (`kernel.batch_shape`, `index_points`, etc.).
        observations_is_missing:  `bool` `Tensor` of shape `[..., e, t]`,
          representing a batch of boolean masks.  When
          `observations_is_missing` is not `None`, this distribution is
          conditioned only on the observations for which the
          corresponding elements of `observations_is_missing` are `False`.
        index_points: (Nested) `float` `Tensor` representing finite collection, or
          batch of collections, of points in the index set over which the GP is
          defined.  Shape (or shape of each nested component) has the form
          `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
          dimensions and must equal `kernel.feature_ndims` (or its corresponding
          nested component) and `e` is the number (size) of index points in each
          batch. Ultimately this distribution corresponds to an `e`-dimensional
          multivariate normal. The batch shape must be broadcastable with
          `kernel.batch_shape`.
        mean_fn: Python `callable` that acts on `index_points` to produce a (batch
          of) collection of mean values at `index_points`. Takes a (nested)
          `Tensor` of shape `[b1, ..., bB, e, f1, ..., fF]` and returns a `Tensor`
          whose shape is broadcastable with `[b1, ..., bB, e, t]`, where `t` is
          the number of tasks.
        observation_noise_variance: `float` `Tensor` representing the variance of
          the noise in the Normal likelihood distribution of the model. May be
          batched, in which case the batch shape must be broadcastable with the
          shapes of all other batched parameters (`kernel.batch_shape`,
          `index_points`, etc.).
          Default value: `None`
        predictive_noise_variance: `float` `Tensor` representing the variance in
          the posterior predictive model. If `None`, we simply re-use
          `observation_noise_variance` for the posterior predictive noise. If set
          explicitly, however, we use this value. This allows us, for example, to
          omit predictive noise variance (by setting this to zero) to obtain
          noiseless posterior predictions of function values, conditioned on noisy
          observations.
        cholesky_fn: Callable which takes a single (batch) matrix argument and
          returns a Cholesky-like lower triangular factor.  Default value: `None`,
            in which case `make_cholesky_with_jitter_fn(1e-6)` is used.
        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.
          Default value: `False`.
        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.
          Default value: `False`.
        name: Python `str` name prefixed to Ops created by this class.
          Default value: 'MultiTaskGaussianProcessRegressionModel'.
        _flattened_conditional_mean_fn: Internal parameter -- do not use.
        _observation_scale: Internal parameter -- do not use.
      """
      parameters = dict(locals())
      with tf.name_scope(name) as name:
          if not isinstance(kernel, multitask_kernel.MultiTaskKernel):
              raise ValueError("`kernel` must be a `MultiTaskKernel`.")
  
          input_dtype = dtype_util.common_dtype(
              dict(
                  kernel=kernel,
                  index_points=index_points,
                  observation_index_points=observation_index_points,
              ),
              dtype_hint=nest_util.broadcast_structure(
                  kernel.feature_ndims, tf.float32
              ),
          )
  
          # If the input dtype is non-nested float, we infer a single dtype for the
          # input and the float parameters, which is also the dtype of the MTGP's
          # samples, log_prob, etc. If the input dtype is nested (or not float), we
          # do not use it to infer the MTGP's float dtype.
          if not tf.nest.is_nested(
              input_dtype
          ) and dtype_util.is_floating(input_dtype):
              dtype = dtype_util.common_dtype(
                  dict(
                      kernel=kernel,
                      index_points=index_points,
                      observations=observations,
                      observation_index_points=observation_index_points,
                      observation_noise_variance=observation_noise_variance,
                      predictive_noise_variance=predictive_noise_variance,
                  ),
                  dtype_hint=tf.float32,
              )
              input_dtype = dtype
          else:
              dtype = dtype_util.common_dtype(
                  dict(
                      observations=observations,
                      observation_noise_variance=observation_noise_variance,
                      predictive_noise_variance=predictive_noise_variance,
                  ),
                  dtype_hint=tf.float32,
              )
  
          if index_points is not None:
              index_points = nest_util.convert_to_nested_tensor(
                  index_points,
                  dtype=input_dtype,
                  convert_ref=False,
                  name="index_points",
                  allow_packing=True,
              )
          if observation_index_points is not None:
              observation_index_points = (
                  nest_util.convert_to_nested_tensor(
                      observation_index_points,
                      dtype=input_dtype,
                      convert_ref=False,
                      name="observation_index_points",
                      allow_packing=True,
                  )
              )
          observations = tensor_util.convert_nonref_to_tensor(
              observations, dtype=dtype, name="observations"
          )
          if observations_is_missing is not None:
              observations_is_missing = (
                  tensor_util.convert_nonref_to_tensor(
                      observations_is_missing, dtype=tf.bool
                  )
              )
          if observation_noise_variance is not None:
              observation_noise_variance = (
                  tensor_util.convert_nonref_to_tensor(
                      observation_noise_variance,
                      dtype=dtype,
                      name="observation_noise_variance",
                  )
              )
          predictive_noise_variance = (
              tensor_util.convert_nonref_to_tensor(
                  predictive_noise_variance,
                  dtype=dtype,
                  name="predictive_noise_variance",
              )
          )
          if predictive_noise_variance is None:
              predictive_noise_variance = observation_noise_variance
          if cholesky_fn is None:
              self._cholesky_fn = (
                  cholesky_util.make_cholesky_with_jitter_fn()
              )
          else:
              if not callable(cholesky_fn):
                  raise ValueError(
                      "`cholesky_fn` must be a Python callable"
                  )
              self._cholesky_fn = cholesky_fn
  
          self._kernel = kernel
          self._index_points = index_points
  
          mean_fn = (
              stochastic_process_util.maybe_create_multitask_mean_fn(
                  mean_fn, kernel, dtype
              )
          )
          self._mean_fn = mean_fn
          self._observation_noise_variance = observation_noise_variance
          self._predictive_noise_variance = predictive_noise_variance
          self._index_ponts = index_points
          self._observation_index_points = observation_index_points
          self._observations = observations
          self._observations_is_missing = observations_is_missing
  
          self._check_observations_valid(observations)
  
          if _flattened_conditional_mean_fn is None:
  
              def flattened_conditional_mean_fn(x):
                  """Flattened Conditional mean."""
                  observation_scale = _compute_observation_scale(
                      kernel,
                      observation_index_points,
                      self._cholesky_fn,
                      observation_noise_variance=self.observation_noise_variance,
                      observations_is_missing=observations_is_missing,
                  )
  
                  return _flattened_conditional_mean_fn_helper(
                      x,
                      self.kernel,
                      self._observations,
                      self._observation_index_points,
                      observations_is_missing,
                      observation_scale,
                      mean_fn,
                  )
  
              _flattened_conditional_mean_fn = (
                  flattened_conditional_mean_fn
              )
  
          self._flattened_conditional_mean_fn = (
              _flattened_conditional_mean_fn
          )
          self._observation_scale = _observation_scale
  
          super(MultiTaskGaussianProcessRegressionModel, self).__init__(
              dtype=dtype,
              reparameterization_type=(
                  reparameterization.FULLY_REPARAMETERIZED
              ),
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              parameters=parameters,
              name=name,
          )
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import cholesky_util
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions.internal import stochastic_process_util
  from tensorflow_probability.python.experimental.distributions import multitask_gaussian_process
  from tensorflow_probability.python.experimental.linalg import linear_operator_unitary
  from tensorflow_probability.python.experimental.psd_kernels import multitask_kernel
  from tensorflow_probability.python.internal import batch_shape_lib
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 1 to create a new MultiTaskGaussianProcessRegressionModel object from inherited base classes: distribution with example: obj = MultiTaskGaussianProcessRegressionModel(kernel, observation_index_points, observations, observations_is_missing, index_points, mean_fn, observation_noise_variance, predictive_noise_variance, cholesky_fn, validate_args, allow_nan_stats, name, _flattened_conditional_mean_fn, _observation_scale)