code: def test_douglas_step_with_batching(self):
      u = np.arange(0, 80, dtype=np.float32).reshape(4, 4, 5)
      d = np.arange(10, 90, dtype=np.float32).reshape(4, 4, 5)
      dx = np.array([d, -3 * d, 2 * d])
      dy = 2 * dx
      dxy = (np.arange(-20, 60, dtype=np.float32).reshape(4, 4, 5)) * 4
      theta = 0.3
      bx = np.arange(20, 100, dtype=np.float32).reshape(4, 4, 5)
      by = np.arange(30, 110, dtype=np.float32).reshape(4, 4, 5)
  
      def equation_params_fn(t):
          del t
          return (
              [
                  [_tfconst(dy), _spread_mixed_term(_tfconst(dxy))],
                  [None, _tfconst(dx)],
              ],
              [_tfconst(by), _tfconst(bx)],
          )
  
      scheme = douglas_adi_scheme(theta=theta)
  
      def pad_fn(t):
          paddings = tf.constant([[0, 0], [0, 1], [1, 0]])
          return tf.pad(t, paddings)
  
      actual = self.evaluate(
          scheme(
              value_grid=tf.constant(u, dtype=tf.float32),
              t1=0,
              t2=1,
              equation_params_fn=equation_params_fn,
              append_boundaries_fn=pad_fn,
              n_dims=2,
              # Tests default boundary as well
              has_default_lower_boundary=[True, False],
              has_default_upper_boundary=[False, True],
          )
      )
      expected = np.zeros_like(u)
      for i in range(4):
          expected[i] = self._simplified_douglas_step_2d(
              u[i],
              dx[:, i],
              dy[:, i],
              dxy[i],
              bx[i],
              by[i],
              0,
              1,
              theta,
          )
  
      self.assertLess(np.max(np.abs(expected - actual)), 0.01)
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/steppers/douglas_adi_scheme_test.py
imports: import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 6 for calling the DouglasAdiSchemeTest obj's underlying member method code with example usage: obj.test_douglas_step_with_batching() and returns: list,  list, tf