code: def test_2d_vector_valued_function(self):
      x_ref_min = np.array([1.0, 0.0], dtype=np.float32)
      x_ref_max = np.array([2.3, 1.0], dtype=np.float32)
      ny = [200, 210]
  
      # Build y_ref.
      x0s, x1s = tf.meshgrid(
          tf.linspace(x_ref_min[0], x_ref_max[0], ny[0]),
          tf.linspace(x_ref_min[1], x_ref_max[1], ny[1]),
          indexing="ij",
      )
  
      def func(x0, x1):
          # Shape [..., 2] output.
          return tf.stack([tf.sin(x0 * x1), tf.cos(x0 * x1)], axis=-1)
  
      # Shape ny + [2]
      y_ref = self.evaluate(func(x0s, x1s))
  
      # Shape [10, 2]
      seed = test_util.test_seed_stream()
      x = tf.stack(
          [
              tf.random.uniform(
                  shape=(10,),
                  minval=x_ref_min[0],
                  maxval=x_ref_max[0],
                  seed=seed(),
              ),
              tf.random.uniform(
                  shape=(10,),
                  minval=x_ref_min[1],
                  maxval=x_ref_max[1],
                  seed=seed(),
              ),
          ],
          axis=-1,
      )
  
      x = self.evaluate(x)
  
      expected_y = func(x[:, 0], x[:, 1])
      actual_y = self.batch_interp_regular_nd_grid_fn(
          x=x,
          x_ref_min=x_ref_min,
          x_ref_max=x_ref_max,
          y_ref=y_ref,
          axis=-3,
      )
  
      self.assertAllClose(
          self.evaluate(expected_y), self.evaluate(actual_y), atol=0.02
      )
  
imports: import numpy
  from scipy import interpolate
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math import interpolation
label_desc: Use image node_id 6 for calling the BaseBatchInterpNDGridTest obj's underlying member method code with example usage: obj.test_2d_vector_valued_function() and returns: tf