docstr: Self-consistency checks with possibly batched inputs.
docstr_tok: ["Self-consistency","checks","with","possibly","batched","inputs","."]
code: def testBatchedInputs(
      self,
      x_batch,
      diag_shift_batch,
      length_scale_batch,
      mat2_batch,
      is_square,
  ):
      """Self-consistency checks with possibly batched inputs."""
      if not self.use_static_shape and tf.executing_eagerly():
          self.skipTest("Cannot test dynamic shape in Eager execution.")
      # Arrange inputs.
      ndims = 2
      x1 = np.stack([np.linspace(-5.0, 5.0, 5)] * ndims, axis=-1)
      x1_length = x1.shape[-2]
      if is_square:
          x2 = None
          x2_length = x1.shape[-2]
      else:
          x2 = np.stack([np.linspace(-5.0, 5.0, 3)] * ndims, axis=-1)
          x2_length = x2.shape[-2]
      mat2 = tf.random.normal(
          mat2_batch + (x2_length, 4),
          dtype=self.dtype,
          seed=test_util.test_seed(),
      )
      bounds_min = [-5.0] * ndims
      bounds_max = [5.0] * ndims
      num_interpolation_points = 7
      diag_shift = 1.0 if is_square else None
      length_scale = 3.0
  
      def _add_batch(v, batch):
          v = np.array(v)
          return np.tile(
              np.reshape(v, (1,) * len(batch) + tuple(v.shape)),
              batch + (1,) * len(v.shape),
          )
  
      if x_batch:
          x1 = _add_batch(x1, x_batch)
      if diag_shift_batch:
          diag_shift = _add_batch(diag_shift, diag_shift_batch)
      if length_scale_batch:
          length_scale = _add_batch(length_scale, length_scale_batch)
  
      # Construct linop, do some operations on it.
      linop = loipk.LinearOperatorInterpolatedPSDKernel(
          kernel=matern.MaternThreeHalves(
              length_scale=self.make_input(length_scale, [])
          ),
          bounds_min=tf.constant(bounds_min, self.dtype),
          bounds_max=tf.constant(bounds_max, self.dtype),
          num_interpolation_points=num_interpolation_points,
          x1=self.make_input(x1, [-2, -1]),
          x2=x2 if x2 is None else self.make_input(x2, [-2, -1]),
          diag_shift=diag_shift
          if diag_shift is None
          else self.make_input(diag_shift, []),
      )
  
      dense = linop.to_dense()
      rows = linop.row([0, 1, 2])
      if is_square:
          diag_part = linop.diag_part()
      matmul = linop.matmul(self.make_input(mat2, [-2, -1]))
      dense_matmul = dense @ mat2
  
      # Perform the tests.
      expected_batch = functools.reduce(
          tf.broadcast_static_shape,
          map(
              tf.TensorShape,
              (x_batch, diag_shift_batch, length_scale_batch),
          ),
      )
      expected_mat_shape = tuple(expected_batch) + (
          x1_length,
          x2_length,
      )
      expected_diag_shape = tuple(expected_batch) + (x1_length,)
      expected_row_shape = tuple(expected_batch) + (3,) + (x2_length,)
      expected_matmul_shape = tuple(
          tf.broadcast_static_shape(
              expected_batch, tf.TensorShape(mat2_batch)
          )
      ) + (x1_length, 4)
  
      with self.subTest("shapes"):
          if self.use_static_shape:
              self.assertAllEqual(expected_mat_shape, dense.shape)
              self.assertAllEqual(expected_mat_shape, linop.shape)
              if is_square:
                  self.assertAllEqual(
                      expected_diag_shape, diag_part.shape
                  )
              self.assertAllEqual(expected_row_shape, rows.shape)
              self.assertAllEqual(expected_matmul_shape, matmul.shape)
  
          self.assertAllEqual(
              expected_mat_shape, self.evaluate(tf.shape(dense))
          )
          self.assertAllEqual(
              expected_mat_shape, self.evaluate(linop.shape_tensor())
          )
          if is_square:
              self.assertAllEqual(
                  expected_diag_shape,
                  self.evaluate(tf.shape(diag_part)),
              )
          self.assertAllEqual(
              expected_row_shape, self.evaluate(tf.shape(rows))
          )
          self.assertAllEqual(
              expected_matmul_shape, self.evaluate(tf.shape(matmul))
          )
  
      with self.subTest("dtypes"):
          self.assertEqual(self.dtype, dense.dtype)
          self.assertEqual(self.dtype, linop.dtype)
          if is_square:
              self.assertEqual(self.dtype, diag_part.dtype)
          self.assertEqual(self.dtype, rows.dtype)
          self.assertEqual(self.dtype, matmul.dtype)
  
      with self.subTest("self_consistency"):
          dense_ = self.evaluate(dense)
          self.assertAllClose(rows, dense_[..., [0, 1, 2], :])
          if is_square:
              expected_diag_part_ = self.evaluate(
                  tf.linalg.diag_part(dense)
              )
              self.assertAllClose(diag_part, expected_diag_part_)
          self.assertAllClose(dense_matmul, matmul)
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.linalg import linear_operator_interpolated_psd_kernel
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import matern
  from tensorflow_probability.python.math.psd_kernels import polynomial
label_desc: Use image node_id 1 for calling the _LinearOperatorInterpolatedPSDKernelTest obj's underlying member method code with example usage: obj.testBatchedInputs(x_batch, diag_shift_batch, length_scale_batch, mat2_batch, is_square) and returns: np