docstr: Performs some minimal self-consistency checks with 3D interpolation.
docstr_tok: ["Performs","some","minimal","self-consistency","checks","with","3D","interpolation","."]
code: def test3D(self, is_square):
      """Performs some minimal self-consistency checks with 3D interpolation."""
      # Arrange inputs.
      ndims = 3
      x1 = np.stack([np.linspace(-5.0, 5.0, 5)] * ndims, axis=-1)
      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(
          (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
  
      # 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.
      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 _LinearOperatorInterpolatedPSDKernelTestStaticTest obj's underlying member method code with example usage: obj.test3D(is_square) without return types