code: def test_backward_pde_coeffs_with_constant_params_2d(self):
      vol = [[1, 2], [3, 4]]
      drift = [1, 2]
      discounting = 3
  
      # batch_shape = (1, ), dim = 2
      # vol_fn(...).shape = batch_shape + (dim, dim) = (1, 2, 2)
      def vol_fn(t, x):
          del t, x
          return tf.expand_dims(tf.constant(vol, dtype=tf.float32), 0)
  
      # drift_fn(...).shape = batch_shape + (dim,) = (1, 2)
      def drift_fn(t, x):
          del t, x
          return tf.expand_dims(tf.constant(drift, dtype=tf.float32), 0)
  
      # discounting_fn(...).shape = batch_shape = (1, )
      discounting_fn = lambda t, x: tf.constant(
          [discounting], dtype=tf.float32
      )
  
      process = tff.models.GenericItoProcess(
          dim=2,
          drift_fn=drift_fn,
          volatility_fn=vol_fn,
          dtype=tf.float32,
      )
  
      pde_solver_fn = mock.Mock()
      coord_grid = [tf.constant([0])]
      process.fd_solver_backward(
          start_time=0,
          end_time=0,
          coord_grid=coord_grid,
          values_grid=tf.constant([0]),
          discounting=discounting_fn,
          pde_solver_fn=pde_solver_fn,
      )
      kwargs = pde_solver_fn.call_args[1]
      second_order_coeff = self.evaluate(
          kwargs["second_order_coeff_fn"](0, coord_grid)
      )
  
      # second_order_coeff.shape = (dim, dim) + grid_shape = (2, 2, 1)
      self.assertAllClose(
          np.array([[[2.5], [5.5]], [[5.5], [12.5]]]),
          second_order_coeff,
      )
  
      # first_order_coeff.shape = (dim) + grid_shape = (2, 1)
      first_order_coeff = self.evaluate(
          kwargs["first_order_coeff_fn"](0, coord_grid)
      )
      self.assertAllClose(np.array([[1], [2]]), first_order_coeff)
  
      # first_order_coeff.shape = grid_shape = (1)
      zeroth_order_coeff = self.evaluate(
          kwargs["zeroth_order_coeff_fn"](0, coord_grid)
      )
      self.assertAllClose(np.array([-3]), zeroth_order_coeff)
  
imports: import os
  from unittest import mock
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 9 for calling the GenericItoProcessTest obj's underlying member method code with example usage: obj.test_backward_pde_coeffs_with_constant_params_2d() and returns: tf, tf