code: def test_backward_pde_coeffs_with_nonconstant_params_2d(self):
      # vol = [[x * y, 2 * x], [x + y, 3]]
      # drift = [x + y, x - y]
      # discounting = 3 * x * y
      # x = [0, 1, 2], y = [-1, 0, 1]
  
      # batch_shape = (3, 3 ), dim = 2
      # vol_fn(...).shape = batch_shape + (dim, dim) = (3, 3, 2, 2)
      def vol_fn(t, grid):
          del t
          self.assertEqual((3, 3, 2), grid.shape)
          ys = grid[..., 0]
          xs = grid[..., 1]
          vol_yy = xs * ys
          vol_yx = 2 * xs
          vol_xy = xs + ys
          vol_xx = 3 * tf.ones_like(xs)
          vol = tf.stack(
              (
                  tf.stack((vol_yy, vol_xy), axis=-1),
                  tf.stack((vol_yx, vol_xx), axis=-1),
              ),
              axis=-1,
          )
          self.assertEqual((3, 3, 2, 2), vol.shape)
          return vol
  
      # drift_fn(...).shape = batch_shape + (dim, ) = (3, 3, 2)
      def drift_fn(t, grid):
          del t
          self.assertEqual((3, 3, 2), grid.shape)
          ys = grid[..., 0]
          xs = grid[..., 1]
          drift_y = xs + ys
          drift_x = xs - ys
          drift = tf.stack((drift_y, drift_x), axis=-1)
          self.assertEqual((3, 3, 2), drift.shape)
          return drift
  
      # drift_fn(...).shape = batch_shape = (3, 3)
      def discounting_fn(t, grid):
          del t
          self.assertEqual((3, 3, 2), grid.shape)
          ys = grid[..., 0]
          xs = grid[..., 1]
          return 3 * xs * ys
  
      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([-1, 0, 1]), tf.constant([0, 1, 2])]
      process.fd_solver_backward(
          start_time=0,
          end_time=0,
          coord_grid=coord_grid,
          values_grid=tf.zeros((3, 3)),
          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)
      )
  
      expected_coeff_yy = np.array(
          [[0, 2.5, 10], [0, 2, 8], [0, 2.5, 10]]
      )
      self.assertAllClose(expected_coeff_yy, second_order_coeff[0][0])
  
      expected_coeff_xy = np.array([[0, 3, 5], [0, 3, 6], [0, 4, 9]])
      self.assertAllClose(expected_coeff_xy, second_order_coeff[0][1])
      self.assertAllClose(expected_coeff_xy, second_order_coeff[1][0])
  
      expected_coeff_xx = np.array(
          [[5, 4.5, 5], [4.5, 5, 6.5], [5, 6.5, 9]]
      )
      self.assertAllClose(expected_coeff_xx, second_order_coeff[1][1])
  
      # first_order_coeff.shape = (dim) + grid_shape = (2, 1)
      first_order_coeff = self.evaluate(
          kwargs["first_order_coeff_fn"](0, coord_grid)
      )
  
      expected_coeff_y = np.array([[-1, 0, 1], [0, 1, 2], [1, 2, 3]])
      self.assertAllClose(expected_coeff_y, first_order_coeff[0])
      expected_coeff_x = np.array([[1, 2, 3], [0, 1, 2], [-1, 0, 1]])
      self.assertAllClose(expected_coeff_x, first_order_coeff[1])
  
      # first_order_coeff.shape = grid_shape = (1)
      zeroth_order_coeff = self.evaluate(
          kwargs["zeroth_order_coeff_fn"](0, coord_grid)
      )
      self.assertAllClose(
          np.array([[0, 3, 6], [0, 0, 0], [0, -3, -6]]),
          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 10 for calling the GenericItoProcessTest obj's underlying member method code with example usage: obj.test_backward_pde_coeffs_with_nonconstant_params_2d() and returns: vol, drift, unknown