code: def testMasking(self):
      seed1, seed2, seed3 = samplers.split_seed(
          test_util.test_seed(), n=3
      )
      base_kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          5.0, 1.0, feature_ndims=1
      )
      mask = np.array(
          [
              [True, False, True, False, True],
              [False, True, False, True, True],
              [False, True, True, False, False],
          ]
      )
      fixed_inputs = tf.where(
          mask[..., np.newaxis],
          tf.random.stateless_uniform(
              [3, 5, 4], seed=seed1, minval=-1.0, maxval=1.0
          ),
          np.nan,
      )
      k = schur_complement.SchurComplement(
          base_kernel, fixed_inputs, fixed_inputs_is_missing=~mask
      )
  
      x = tf.random.stateless_uniform(
          [6, 4], seed=seed2, minval=-1.0, maxval=1.0
      )
      y = tf.random.stateless_uniform(
          [2, 1, 2, 4], seed=seed3, minval=-1.0, maxval=1.0
      )
  
      k_x_y = k.matrix(x, y)
      k_x_x = k.apply(x, x, example_ndims=1)
      k_y_y = k.apply(y, y, example_ndims=1)
  
      # For each batch member of `k`, compare `k.matrix` and `k.apply` to the
      # Schur complement computed using only the subset `fixed_inputs` that are
      # not masked out.
      for i in range(3):
          fixed_i = tf.gather(fixed_inputs[i], mask[i].nonzero()[0])
          k_i = schur_complement.SchurComplement(base_kernel, fixed_i)
          self.assertAllClose(
              k_i.matrix(x, y),
              k_x_y[:, i : i + 1],
              atol=1e-5,
              rtol=1e-5,
          )
          self.assertAllClose(
              k_i.apply(x, x), k_x_x[i], atol=1e-5, rtol=1e-5
          )
          self.assertAllClose(
              k_i.apply(y, y), k_y_y[:, i : i + 1], atol=1e-5, rtol=1e-5
          )
  
      def _grad(x):
          return gradient.value_and_gradient(
              lambda u: tf.reduce_sum(k.apply(u, u, example_ndims=1)), x
          )[1]
  
      self.assertAllNotNan(_grad(x))
  
imports: import functools
  import itertools
  from unittest import mock
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import cholesky_util
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.math.psd_kernels import schur_complement
  from tensorflow_probability.python.math.psd_kernels.internal import test_util
label_desc: Use image node_id 3 for calling the SchurComplementTest obj's underlying member method code with example usage: obj.testMasking() and returns: gradient