code: def testBatchShape(self):
      # 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)
      # Use 3 feature_ndims.
      kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          amplitude, inner_length_scale, feature_ndims=3
      )
      scale_diag = tf.ones([20, 1, 2, 1, 1, 1], dtype=self.dtype)
      ard_kernel = feature_scaled.FeatureScaled(
          kernel, scale_diag=scale_diag
      )
      self.assertAllEqual([20, 10, 2], ard_kernel.batch_shape)
      self.assertAllEqual(
          [20, 10, 2], self.evaluate(ard_kernel.batch_shape_tensor())
      )
  
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 1 for calling the _FeatureScaledTest obj's underlying member method code with example usage: obj.testBatchShape() without return types