code: def testTensorShapesAreCorrect(self):
      for x1_example_ndims in range(0, 3):
          for x2_example_ndims in range(0, 2):
              # An integer generator.
              ints = itertools.count(start=2, step=1)
              feature_shape = [next(ints), next(ints)]
  
              x_batch_shape = [next(ints)]
              y_batch_shape = [next(ints), 1]
              z_batch_shape = [next(ints), 1, 1]
  
              num_x = [next(ints) for _ in range(x1_example_ndims)]
              num_y = [next(ints) for _ in range(x2_example_ndims)]
              num_z = [next(ints)]
  
              x_shape = x_batch_shape + num_x + feature_shape
              y_shape = y_batch_shape + num_y + feature_shape
              z_shape = z_batch_shape + num_z + feature_shape
  
              x = np.ones(x_shape, np.float64)
              y = np.ones(y_shape, np.float64)
              z = np.random.uniform(-1.0, 1.0, size=z_shape)
  
              base_kernel = (
                  exponentiated_quadratic.ExponentiatedQuadratic(
                      amplitude=np.ones(
                          [next(ints), 1, 1, 1], np.float64
                      ),
                      feature_ndims=len(feature_shape),
                  )
              )
  
              k = schur_complement.SchurComplement(
                  base_kernel, fixed_inputs=z
              )
  
              expected = (
                  broadcast_shapes(
                      base_kernel.batch_shape,
                      x_batch_shape,
                      y_batch_shape,
                      z_batch_shape,
                  )
                  + num_x
                  + num_y
              )
  
              mat = k.tensor(
                  x,
                  y,
                  x1_example_ndims=x1_example_ndims,
                  x2_example_ndims=x2_example_ndims,
              )
              actual = mat.shape
              self.assertAllEqual(expected, actual)
  
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 5 for calling the SchurComplementTest obj's underlying member method code with example usage: obj.testTensorShapesAreCorrect() without return types