code: def test_2d_vector_valued_function_with_batch_dims(self):
      # No batch dims, will broadcast.
      x_ref_min = np.array([0.0, 0.0], dtype=np.float32)
  
      # No batch dims, will broadcast.
      x_ref_max = np.array([1.0, 1.0], dtype=np.float32)
      ny = [200, 210]
  
      # Build y_ref.
  
      # First step is to build up two batches of x0 and x1.
      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",
      )
      x0s = tf.stack([x0s, x0s], axis=0)
      x1s = tf.stack([x1s, x1s], axis=0)
  
      def func(batch_of_x0, batch_of_x1):
          """Function that does something different for batch 0 and batch 1."""
          # batch_0_result.shape = [..., 2].
          x0, x1 = batch_of_x0[0, ...], batch_of_x1[0, ...]
          batch_0_result = tf.stack(
              [tf.sin(x0 * x1), tf.cos(x0 * x1)], axis=-1
          )
  
          x0, x1 = batch_of_x0[1, ...], batch_of_x1[1, ...]
          batch_1_result = tf.stack(
              [tf.sin(2 * x0), tf.cos(2 * x1)], axis=-1
          )
  
          return tf.stack([batch_0_result, batch_1_result], axis=0)
  
      # Shape [2] + ny + [2]
      y_ref = self.evaluate(func(x0s, x1s))
  
      # Shape [2, 10, 2].  The batch shape is [2], the [10] is the number of
      # interpolants per batch.
      x = tf.random.uniform(
          shape=[2, 10, 2], seed=test_util.test_seed()
      )
  
      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 7 for calling the BaseBatchInterpNDGridTest obj's underlying member method code with example usage: obj.test_2d_vector_valued_function_with_batch_dims() and returns: tf