class_docstr: The SchurComplement kernel.
  
  Given a block matrix `M = [[A, B], [C, D]]`, the Schur complement of D in M is
  written `M / D = A - B @ Inverse(D) @ C`.
  
  This class represents a PositiveSemidefiniteKernel whose behavior is as
  follows. We compute a matrix, analogous to `D` in the above definition, by
  calling `base_kernel.matrix(fixed_inputs, fixed_inputs)`. Then given new input
  locations `x` and `y`, we can construct the remaining pieces of `M` above, and
  compute the Schur complement of `D` in `M` (see Mathematical Details, below).
  
  Notably, this kernel uses a bijector (Invert(CholeskyOuterProduct)), as an
  intermediary for the requisite matrix solve, which means we get a caching
  benefit after the first use.
  
  ### Mathematical Details
  
  Suppose we have a kernel `k` and some fixed collection of inputs
  `Z = [z0, z1, ..., zN]`. Given new inputs `x` and `y`, we can form a block
  matrix
  
   ```none
     M = [
       [k(x, y), k(x, z0), ..., k(x, zN)],
       [k(z0, y), k(z0, z0), ..., k(z0, zN)],
       ...,
       [k(zN, y), k(z0, zN), ..., k(zN, zN)],
     ]
   ```
  
  We might write this, so as to emphasize the block structure,
  
   ```none
     M = [
       [xy, xZ],
       [yZ^T, ZZ],
     ],
  
     xy = [k(x, y)]
     xZ = [k(x, z0), ..., k(x, zN)]
     yZ = [k(y, z0), ..., k(y, zN)]
     ZZ = "the matrix of k(zi, zj)'s"
   ```
  
  Then we have the definition of this kernel's apply method:
  
  `schur_comp.apply(x, y) = xy - xZ @ ZZ^{-1} @ yZ^T`
  
  and similarly, if x and y are collections of inputs.
  
  As with other PSDKernels, the `apply` method acts as a (possibly
  vectorized) scalar function of 2 inputs. Given a single `x` and `y`,
  `apply` will yield a scalar output. Given two (equal size!) collections `X`
  and `Y`, it will yield another (equal size!) collection of scalar outputs.
  
  ### Examples
  
  Here's a simple example usage, with no particular motivation.
  
  ```python
  from tensorflow_probability.math import psd_kernels
  
  base_kernel = psd_kernels.ExponentiatedQuadratic(amplitude=np.float64(1.))
  # 3 points in 1-dimensional space (shape [3, 1]).
  z = [[0.], [3.], [4.]]
  
  schur_kernel = psd_kernels.SchurComplement(
      base_kernel=base_kernel,
      fixed_inputs=z)
  
  # Two individual 1-d points
  x = [1.]
  y = [2.]
  print(schur_kernel.apply(x, y))
  # ==> k(x, y) - k(x, z) @ Inverse(k(z, z)) @ k(z, y)
  ```
  
  A more motivating application of this kernel is in constructing a Gaussian
  process that is conditioned on some observed data.
  
  ```python
  from tensorflow_probability import distributions as tfd
  from tensorflow_probability.math import psd_kernels
  
  base_kernel = psd_kernels.ExponentiatedQuadratic(amplitude=np.float64(1.))
  observation_index_points = np.random.uniform(-1., 1., [50, 1])
  observations = np.sin(2 * np.pi * observation_index_points[..., 0])
  
  posterior_kernel = psd_kernels.SchurComplement(
      base_kernel=base_kernel,
      fixed_inputs=observation_index_points)
  
  # Assume we use a zero prior mean, and compute the posterior mean.
  def posterior_mean_fn(x):
    k_x_obs_linop = tf.linalg.LinearOperatorFullMatrix(
        base_kernel.matrix(x, observation_index_points))
    chol_linop = tf.linalg.LinearOperatorLowerTriangular(
        posterior_kernel.divisor_matrix_cholesky())
  
    return k_x_obs_linop.matvec(
        chol_linop.solvevec(
            chol_linop.solvevec(observations),
            adjoint=True))
  
  # Construct the GP posterior distribution at some new points.
  gp_posterior = tfp.distributions.GaussianProcess(
      index_points=np.linspace(-1., 1., 100)[..., np.newaxis],
      kernel=posterior_kernel,
      mean_fn=posterior_mean_fn)
  
  # Draw 5 samples on the above 100-point grid
  samples = gp_posterior.sample(5)
  ```
docstr: Returns a `SchurComplement` with a precomputed divisor matrix.
  
  This method is the same as creating a `SchurComplement` kernel, but assumes
  that `fixed_inputs`, `diag_shift` and `base_kernel` are unchanging /
  not parameterized by any mutable state. We explicitly read / concretize
  these values when this method is called, since we can precompute some
  factorizations in order to speed up subsequent invocations of the kernel.
  
  WARNING: This method assumes passed in arguments are not parameterized
  by mutable state (`fixed_inputs`, `diag_shift` and `base_kernel`), and hence
  is not tape-safe.
  
  Args:
    base_kernel: A `PositiveSemidefiniteKernel` instance, the kernel used to
      build the block matrices of which this kernel computes the Schur
      complement.
    fixed_inputs: A (nested) Tensor, representing a collection of inputs. The
      Schur complement that this kernel computes comes from a block matrix,
      whose bottom-right corner is derived from
      `base_kernel.matrix(fixed_inputs, fixed_inputs)`, and whose top-right
      and bottom-left pieces are constructed by computing the base_kernel
      at pairs of input locations together with these `fixed_inputs`.
      `fixed_inputs` is allowed to be an empty collection (either `None` or
      having a zero shape entry), in which case the kernel falls back to
      the trivial application of `base_kernel` to inputs. See class-level
      docstring for more details on the exact computation this does;
      `fixed_inputs` correspond to the `Z` structure discussed there.
      `fixed_inputs` (or each of its nested components) is assumed to have
      shape `[b1, ..., bB, N, f1, ..., fF]` where the `b`'s are batch shape
      entries, the `f`'s are feature_shape entries, and `N` is the number
      of fixed inputs. Use of this kernel entails a 1-time O(N^3) cost of
      computing the Cholesky decomposition of the k(Z, Z) matrix. The batch
      shape elements of `fixed_inputs` must be broadcast compatible with
      `base_kernel.batch_shape`.
    fixed_inputs_is_missing: A boolean Tensor of shape `[..., N]`.  When
      `is_missing` is not None and an element of `is_missing` is `True`, the
      returned kernel will return values computed as if the divisor matrix
      did not contain the corresponding row or column.
    diag_shift: A floating point scalar to be added to the diagonal of the
      divisor_matrix before computing its Cholesky.
    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` is used with the `jitter`
      parameter.
    validate_args: If `True`, parameters are checked for validity despite
      possibly degrading runtime performance.
      Default value: `False`
    name: Python `str` name prefixed to Ops created by this class.
      Default value: `"PrecomputedSchurComplement"`
    _precomputed_divisor_matrix_cholesky: Internal arg -- do not use.
docstr_tok: ["Returns","a","`","SchurComplement","`","with","a","precomputed","divisor","matrix",".","This","method","is","the","same","as","creating","a","`","SchurComplement","`","kernel",",","but","assumes","that","`","fixed_inputs","`",",","`","diag_shift","`","and","`","base_kernel","`","are","unchanging","\/","not","parameterized","by","any","mutable","state",".","We","explicitly","read","\/","concretize","these","values","when","this","method","is","called",",","since","we","can","precompute","some","factorizations","in","order","to","speed","up","subsequent","invocations","of","the","kernel",".","WARNING",":","This","method","assumes","passed","in","arguments","are","not","parameterized","by","mutable","state","(","`","fixed_inputs","`",",","`","diag_shift","`","and","`","base_kernel","`",")",",","and","hence","is","not","tape-safe",".","Args",":","base_kernel",":","A","`","PositiveSemidefiniteKernel","`","instance",",","the","kernel","used","to","build","the","block","matrices","of","which","this","kernel","computes","the","Schur","complement",".","fixed_inputs",":","A","(","nested",")","Tensor",",","representing","a","collection","of","inputs",".","The","Schur","complement","that","this","kernel","computes","comes","from","a","block","matrix",",","whose","bottom-right","corner","is","derived","from","`","base_kernel.matrix","(","fixed_inputs",",","fixed_inputs",")","`",",","and","whose","top-right","and","bottom-left","pieces","are","constructed","by","computing","the","base_kernel","at","pairs","of","input","locations","together","with","these","`","fixed_inputs","`",".","`","fixed_inputs","`","is","allowed","to","be","an","empty","collection","(","either","`","None","`","or","having","a","zero","shape","entry",")",",","in","which","case","the","kernel","falls","back","to","the","trivial","application","of","`","base_kernel","`","to","inputs",".","See","class-level","docstring","for","more","details","on","the","exact","computation","this","does",";","`","fixed_inputs","`","correspond","to","the","`","Z","`","structure","discussed","there",".","`","fixed_inputs","`","(","or","each","of","its","nested","components",")","is","assumed","to","have","shape","`","[","b1",",","...",",","bB",",","N",",","f1",",","...",",","fF","]","`","where","the","`","b","`","'s","are","batch","shape","entries",",","the","`","f","`","'s","are","feature_shape","entries",",","and","`","N","`","is","the","number","of","fixed","inputs",".","Use","of","this","kernel","entails","a","1-time","O","(","N^3",")","cost","of","computing","the","Cholesky","decomposition","of","the","k","(","Z",",","Z",")","matrix",".","The","batch","shape","elements","of","`","fixed_inputs","`","must","be","broadcast","compatible","with","`","base_kernel.batch_shape","`",".","fixed_inputs_is_missing",":","A","boolean","Tensor","of","shape","`","[","...",",","N","]","`",".","When","`","is_missing","`","is","not","None","and","an","element","of","`","is_missing","`","is","`","True","`",",","the","returned","kernel","will","return","values","computed","as","if","the","divisor","matrix","did","not","contain","the","corresponding","row","or","column",".","diag_shift",":","A","floating","point","scalar","to","be","added","to","the","diagonal","of","the","divisor_matrix","before","computing","its","Cholesky",".","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","`","is","used","with","the","`","jitter","`","parameter",".","validate_args",":","If","`","True","`",",","parameters","are","checked","for","validity","despite","possibly","degrading","runtime","performance",".","Default","value",":","`","False","`","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","Default","value",":","`","``","PrecomputedSchurComplement","''","`","_precomputed_divisor_matrix_cholesky",":","Internal","arg","--","do","not","use","."]
code: def with_precomputed_divisor(
      base_kernel,
      fixed_inputs,
      fixed_inputs_is_missing=None,
      diag_shift=None,
      cholesky_fn=None,
      validate_args=False,
      name="PrecomputedSchurComplement",
      _precomputed_divisor_matrix_cholesky=None,
  ):
      """Returns a `SchurComplement` with a precomputed divisor matrix.
  
      This method is the same as creating a `SchurComplement` kernel, but assumes
      that `fixed_inputs`, `diag_shift` and `base_kernel` are unchanging /
      not parameterized by any mutable state. We explicitly read / concretize
      these values when this method is called, since we can precompute some
      factorizations in order to speed up subsequent invocations of the kernel.
  
      WARNING: This method assumes passed in arguments are not parameterized
      by mutable state (`fixed_inputs`, `diag_shift` and `base_kernel`), and hence
      is not tape-safe.
  
      Args:
        base_kernel: A `PositiveSemidefiniteKernel` instance, the kernel used to
          build the block matrices of which this kernel computes the Schur
          complement.
        fixed_inputs: A (nested) Tensor, representing a collection of inputs. The
          Schur complement that this kernel computes comes from a block matrix,
          whose bottom-right corner is derived from
          `base_kernel.matrix(fixed_inputs, fixed_inputs)`, and whose top-right
          and bottom-left pieces are constructed by computing the base_kernel
          at pairs of input locations together with these `fixed_inputs`.
          `fixed_inputs` is allowed to be an empty collection (either `None` or
          having a zero shape entry), in which case the kernel falls back to
          the trivial application of `base_kernel` to inputs. See class-level
          docstring for more details on the exact computation this does;
          `fixed_inputs` correspond to the `Z` structure discussed there.
          `fixed_inputs` (or each of its nested components) is assumed to have
          shape `[b1, ..., bB, N, f1, ..., fF]` where the `b`'s are batch shape
          entries, the `f`'s are feature_shape entries, and `N` is the number
          of fixed inputs. Use of this kernel entails a 1-time O(N^3) cost of
          computing the Cholesky decomposition of the k(Z, Z) matrix. The batch
          shape elements of `fixed_inputs` must be broadcast compatible with
          `base_kernel.batch_shape`.
        fixed_inputs_is_missing: A boolean Tensor of shape `[..., N]`.  When
          `is_missing` is not None and an element of `is_missing` is `True`, the
          returned kernel will return values computed as if the divisor matrix
          did not contain the corresponding row or column.
        diag_shift: A floating point scalar to be added to the diagonal of the
          divisor_matrix before computing its Cholesky.
        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` is used with the `jitter`
          parameter.
        validate_args: If `True`, parameters are checked for validity despite
          possibly degrading runtime performance.
          Default value: `False`
        name: Python `str` name prefixed to Ops created by this class.
          Default value: `"PrecomputedSchurComplement"`
        _precomputed_divisor_matrix_cholesky: Internal arg -- do not use.
      """
      if tf.nest.is_nested(base_kernel.feature_ndims):
          dtype = dtype_util.common_dtype(
              [base_kernel, fixed_inputs],
              dtype_hint=nest_util.broadcast_structure(
                  base_kernel.feature_ndims, tf.float32
              ),
          )
          float_dtype = dtype_util.common_dtype(
              [diag_shift], tf.float32
          )
      else:
          # If the fixed inputs are not nested, we assume they are of the same
          # float dtype as the remaining parameters.
          dtype = dtype_util.common_dtype(
              [base_kernel, fixed_inputs, diag_shift], tf.float32
          )
          float_dtype = dtype
      fixed_inputs = nest_util.convert_to_nested_tensor(
          fixed_inputs, dtype=dtype, allow_packing=True
      )
      if fixed_inputs_is_missing is not None:
          fixed_inputs_is_missing = tf.convert_to_tensor(
              fixed_inputs_is_missing, tf.bool
          )
      if diag_shift is not None:
          diag_shift = tf.convert_to_tensor(diag_shift, float_dtype)
  
      if cholesky_fn is None:
          from tensorflow_probability.python.distributions import (
              cholesky_util,
          )  # pylint:disable=g-import-not-at-top
  
          cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn()
  
      divisor_matrix_cholesky = _precomputed_divisor_matrix_cholesky
      if divisor_matrix_cholesky is None:
          # TODO(b/196219597): Add a check to ensure that we have a `base_kernel`
          # that is explicitly concretized.
          divisor_matrix_cholesky = cholesky_fn(
              _compute_divisor_matrix(
                  base_kernel,
                  diag_shift=diag_shift,
                  fixed_inputs=fixed_inputs,
                  fixed_inputs_is_missing=fixed_inputs_is_missing,
              )
          )
  
      schur_complement = SchurComplement(
          base_kernel=base_kernel,
          fixed_inputs=fixed_inputs,
          fixed_inputs_is_missing=fixed_inputs_is_missing,
          diag_shift=diag_shift,
          cholesky_fn=cholesky_fn,
          validate_args=validate_args,
          _precomputed_divisor_matrix_cholesky=divisor_matrix_cholesky,
          name=name,
      )
  
      return schur_complement
  
imports: import tensorflow.compat.v2
  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 tensor_util
  from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel
  from tensorflow_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 2 for calling the SchurComplement obj's underlying member method code with example usage: obj.with_precomputed_divisor(base_kernel, fixed_inputs, fixed_inputs_is_missing, diag_shift, cholesky_fn, validate_args, name, _precomputed_divisor_matrix_cholesky) and returns: schur_complement