code: def testShapesAreCorrect(self):
      k = spectral_mixture.SpectralMixture(
          logits=np.ones([6, 7], np.float32),
          locs=np.ones([9, 1, 1, 7, 3], np.float32),
          scales=np.ones([2, 1, 7, 3], np.float32),
      )
  
      batch_shape = [9, 2, 6]
  
      x = np.ones([4, 3], np.float32)
      y = np.ones([5, 3], np.float32)
  
      self.assertAllEqual(batch_shape + [4, 5], k.matrix(x, y).shape)
  
      x = np.ones([2, 1, 3], np.float32)
      y = np.ones([1, 4, 3], np.float32)
      self.assertAllEqual(
          batch_shape + [2, 4], k.apply(x, y, example_ndims=2).shape
      )
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/spectral_mixture_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import spectral_mixture
label_desc: Use image node_id 3 for calling the SpectralMixtureTest obj's underlying member method code with example usage: obj.testShapesAreCorrect() without return types