code: def testSeparableShape(self):
      amplitude = np.random.uniform(2, 3.0, size=[3, 1, 2]).astype(
          self.dtype
      )
      length_scale = np.random.uniform(2, 3.0, size=[1, 3, 1]).astype(
          self.dtype
      )
      base_kernel = exponentiated_quadratic.ExponentiatedQuadratic(
          amplitude, length_scale
      )
      task_kernel_matrix_linop = tf.linalg.LinearOperatorIdentity(
          5, batch_shape=[4, 1, 1, 1], dtype=self.dtype
      )
  
      kernel = multitask_kernel.Separable(
          num_tasks=5,
          base_kernel=base_kernel,
          task_kernel_matrix_linop=task_kernel_matrix_linop,
      )
      self.assertAllEqual(
          [4, 3, 3, 2], self.evaluate(kernel.batch_shape_tensor())
      )
      matrix_over_all_tasks = kernel.matrix_over_all_tasks(
          np.ones([4, 1, 1, 1, 3, 2], dtype=self.dtype),
          np.ones([5, 1, 1, 1, 1, 4, 2], dtype=self.dtype),
      )
      self.assertAllEqual(
          [5, 4, 3, 3, 2, 15, 20],
          self.evaluate(matrix_over_all_tasks.shape_tensor()),
      )
  
file_path: probability/tensorflow_probability/python/experimental/psd_kernels/multitask_kernel_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.psd_kernels import multitask_kernel
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
label_desc: Use image node_id 2 for calling the _MultiTaskKernelTest obj's underlying member method code with example usage: obj.testSeparableShape() without return types