code: def testKernelParametersBroadcast(self, feature_ndims, dims):
      # Batch shape [10, 2]
      amplitude = np.random.uniform(
          low=1.0, high=10.0, size=[10, 2]
      ).astype(self.dtype)
      inner_length_scale = self.dtype(1.0)
      kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          amplitude, inner_length_scale, feature_ndims=feature_ndims
      )
      input_shape = [dims] * feature_ndims
  
      # Batch shape [3, 1, 2].
      concentration1 = np.random.uniform(
          2, 5, size=([3, 1, 2] + input_shape)
      ).astype(self.dtype)
  
      concentration0 = np.random.uniform(
          2, 5, size=([3, 1, 2] + input_shape)
      ).astype(self.dtype)
  
      kum_kernel = kumaraswamy_transformed.KumaraswamyTransformed(
          kernel, concentration1, concentration0
      )
  
      x = np.random.uniform(size=input_shape).astype(self.dtype)
      y = np.random.uniform(size=input_shape).astype(self.dtype)
  
      self.assertAllClose(
          _numpy_exp_quad(
              amplitude,
              inner_length_scale,
              _kumaraswamy_warp(x, concentration1, concentration0),
              _kumaraswamy_warp(y, concentration1, concentration0),
              feature_ndims=feature_ndims,
          ),
          self.evaluate(kum_kernel.apply(x, y)),
          rtol=1e-4,
          atol=1e-4,
      )
  
      z = np.random.uniform(size=[10] + input_shape).astype(self.dtype)
  
      expanded_c1 = np.expand_dims(concentration1, -(feature_ndims + 1))
      expanded_c0 = np.expand_dims(concentration0, -(feature_ndims + 1))
  
      self.assertAllClose(
          _numpy_exp_quad_matrix(
              amplitude[..., None, None],
              inner_length_scale,
              _kumaraswamy_warp(z, expanded_c1, expanded_c0),
              feature_ndims=feature_ndims,
          ),
          self.evaluate(kum_kernel.matrix(z, z)),
          atol=1e-4,
          rtol=1e-4,
      )
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/kumaraswamy_transformed_test.py
imports: from absl.testing import parameterized
  import numpy
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.math.psd_kernels import kumaraswamy_transformed
label_desc: Use image node_id 2 for calling the _KumaraswamyTransformedTest obj's underlying member method code with example usage: obj.testKernelParametersBroadcast(feature_ndims, dims) without return types