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,
          length_scale=inner_length_scale,
          feature_ndims=feature_ndims,
      )
      input_shape = [dims] * feature_ndims
  
      # Batch shape [3, 1, 2].
      length_scale = np.random.uniform(
          2, 5, size=([3, 1, 2] + input_shape)
      ).astype(self.dtype)
  
      ard_kernel = feature_scaled.FeatureScaled(
          kernel, scale_diag=length_scale
      )
      self.assertIs(
          dtype_util.as_numpy_dtype(ard_kernel.dtype), self.dtype
      )
  
      x = np.random.uniform(-1, 1, size=input_shape).astype(self.dtype)
      y = np.random.uniform(-1, 1, size=input_shape).astype(self.dtype)
  
      self.assertAllClose(
          _numpy_exp_quad(
              amplitude,
              inner_length_scale,
              x / length_scale,
              y / length_scale,
              feature_ndims=feature_ndims,
          ),
          self.evaluate(ard_kernel.apply(x, y)),
      )
  
      z = np.random.uniform(-1, 1, size=[10] + input_shape).astype(
          self.dtype
      )
  
      expanded_length_scale = np.expand_dims(
          length_scale, -(feature_ndims + 1)
      )
  
      self.assertAllClose(
          _numpy_exp_quad_matrix(
              amplitude[..., None, None],
              inner_length_scale,
              z / expanded_length_scale,
              feature_ndims=feature_ndims,
          ),
          self.evaluate(ard_kernel.matrix(z, z)),
      )
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/feature_scaled_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  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 feature_scaled
label_desc: Use image node_id 2 for calling the _FeatureScaledTest obj's underlying member method code with example usage: obj.testKernelParametersBroadcast(feature_ndims, dims) without return types