code: def testStructuredIndexPoints(
      self, foo_feature_shape, bar_feature_shape
  ):
      base_kernel = psd_kernels.ExponentiatedQuadratic()
      structured_kernel = psd_kernel_test_util.MultipartTestKernel(
          base_kernel,
          feature_ndims={
              "foo": len(foo_feature_shape),
              "bar": len(bar_feature_shape),
          },
      )
  
      foo_num_features = np.prod(foo_feature_shape)
      bar_num_features = np.prod(bar_feature_shape)
      batch_and_example_shape = [3, 2, 10]
      index_points = np.random.uniform(
          -1,
          1,
          batch_and_example_shape
          + [foo_num_features + bar_num_features],
      ).astype(np.float32)
      split_index_points = tf.split(
          index_points, [foo_num_features, bar_num_features], axis=-1
      )
      structured_index_points = {
          "foo": tf.reshape(
              split_index_points[0],
              batch_and_example_shape + foo_feature_shape,
          ),
          "bar": tf.reshape(
              split_index_points[1],
              batch_and_example_shape + bar_feature_shape,
          ),
      }
      base_gp = gaussian_process.GaussianProcess(
          base_kernel, index_points=index_points
      )
      structured_gp = gaussian_process.GaussianProcess(
          structured_kernel, index_points=structured_index_points
      )
  
      self.assertAllEqual(
          base_gp.event_shape, structured_gp.event_shape
      )
      self.assertAllEqual(
          base_gp.event_shape_tensor(),
          structured_gp.event_shape_tensor(),
      )
      self.assertAllEqual(
          base_gp.batch_shape, structured_gp.batch_shape
      )
      self.assertAllEqual(
          base_gp.batch_shape_tensor(),
          structured_gp.batch_shape_tensor(),
      )
  
      s = structured_gp.sample(3, seed=test_util.test_seed())
      self.assertAllClose(
          base_gp.log_prob(s), structured_gp.log_prob(s)
      )
      self.assertAllClose(base_gp.mean(), structured_gp.mean())
      self.assertAllClose(base_gp.variance(), structured_gp.variance())
  
      # Check that batch shapes and number of index points broadcast across
      # different parts of index_points.
      bcast_structured_index_points = {
          "foo": np.random.uniform(
              -1, 1, [2, 1] + foo_feature_shape
          ).astype(np.float32),
          "bar": np.random.uniform(
              -1, 1, [3, 1, 10] + bar_feature_shape
          ).astype(np.float32),
      }
      bcast_structured_gp = gaussian_process.GaussianProcess(
          structured_kernel, index_points=bcast_structured_index_points
      )
      self.assertAllEqual(
          base_gp.event_shape, bcast_structured_gp.event_shape
      )
      self.assertAllEqual(
          base_gp.event_shape_tensor(),
          bcast_structured_gp.event_shape_tensor(),
      )
      self.assertAllEqual(
          base_gp.batch_shape, bcast_structured_gp.batch_shape
      )
      self.assertAllEqual(
          base_gp.batch_shape_tensor(),
          bcast_structured_gp.batch_shape_tensor(),
      )
  
      index_points_bad_num_examples = {
          "foo": np.random.uniform(
              -1, 1, [5] + foo_feature_shape
          ).astype(np.float32),
          "bar": np.random.uniform(
              -1, 1, [2] + bar_feature_shape
          ).astype(np.float32),
      }
  
      # Non-broadcasting numbers of index points should raise.
      structured_gp_bad_num_examples = gaussian_process.GaussianProcess(
          structured_kernel, index_points=index_points_bad_num_examples
      )
      with self.assertRaisesRegex(
          ValueError, "the same or broadcastable"
      ):
          _ = structured_gp_bad_num_examples.event_shape
  
      # Iterable index points should be interpreted as single Tensors if the
      # kernel is not structured.
      index_points_list = tf.unstack(index_points)
      gp_with_list = gaussian_process.GaussianProcess(
          base_kernel, index_points=index_points_list
      )
      self.assertAllEqual(
          base_gp.event_shape_tensor(),
          gp_with_list.event_shape_tensor(),
      )
      self.assertAllEqual(
          base_gp.batch_shape_tensor(),
          gp_with_list.batch_shape_tensor(),
      )
      self.assertAllClose(base_gp.log_prob(s), gp_with_list.log_prob(s))
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gaussian_process
  from tensorflow_probability.python.distributions import gaussian_process_regression_model
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math import psd_kernels
  from tensorflow_probability.python.math.psd_kernels.internal import test_util
label_desc: Use image node_id 14 for calling the _GaussianProcessTest obj's underlying member method code with example usage: obj.testStructuredIndexPoints(foo_feature_shape, bar_feature_shape) without return types