code: def test_backward_pde_coeffs_with_batching_1d(self):
      dtype = tf.float64
  
      # Specify volatilities, interest rates and strikes for the options
      volatilities = tf.constant([[0.3], [0.15], [0.1]], dtype)
      rates = tf.constant([[0.01], [0.03], [0.01]], dtype)
      expiries = 1.0
  
      # Define Generic Ito Process
  
      # Process dimensionality
      dim = 1
  
      # Batch size of the process
      num_processes = 3
  
      def drift_fn(t, x):
          del t
          # `x` is expected to be of shape [num_processes] + sample_shape + [dim]
          # We need to expand rank of rates to
          # `[num_processes] + extra_rank * [1] + [1]`
          expand_rank = x.shape.rank - 2
          rates_expand = tf.reshape(
              rates, [num_processes] + (expand_rank + 1) * [1]
          )
          # Output is of shape [num_processes] + sample_shape + [dim]
          return rates_expand * x
  
      def vol_fn(t, x):
          del t
          # `x` is expeceted to be of shape [num_processes] + sample_shape + [dim]
          # As before, need to expand rank of volatilities to
          # `[num_processes] + extra_rank * [1] + [1]`
          expand_rank = x.shape.rank - 2
          volatilities_expand = tf.reshape(
              volatilities, [num_processes] + (expand_rank + 1) * [1]
          )
          # Output is of shape [num_processes] + sample_shape + [dim, dim]
          return tf.expand_dims(
              volatilities_expand * x, axis=-1
          ) * tf.eye(
              dim, batch_shape=x.shape.as_list()[:-1], dtype=x.dtype
          )
  
      process = tff.models.GenericItoProcess(
          dim=dim, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype
      )
      # Define a 2 strikes for each batch process,
      num_strikes = 2
      # Shape [num_processes, num_strikes, 1]. Here 1 at the end is just for
      # convenience
      strikes = tf.constant(
          [[[50], [60]], [[100], [90]], [[120], [90]]], dtype
      )
  
      # Price a batch of European call options
      @tff.math.pde.boundary_conditions.dirichlet
      def upper_boundary_fn(t, grid):
          del grid
          # Shape (num_processes, num_strikes)
          return tf.squeeze(
              s_max - strikes * tf.exp(-rates * (expiries - t))
          )
  
      # Define discounting function
      def discounting(t, x):
          del t, x
          rates_expand = tf.expand_dims(rates, axis=-1)
          # Shape compatible with (num_processes, num_strikes)
          return rates_expand
  
      # Build a uniform grid
      s_min = 0
      s_max = 200
      num_grid_points = 256  # Number of grid points
  
      grid = tff.math.pde.grids.uniform_grid(
          minimums=[s_min],
          maximums=[s_max],
          sizes=[num_grid_points],
          dtype=dtype,
      )
  
      # Shape [num_processes, num_strikes, num_grid_points]
      final_value_grid = tf.nn.relu(grid[0] - strikes)
  
      prices_estimated = process.fd_solver_backward(
          start_time=expiries,
          end_time=0,
          time_step=0.1,
          coord_grid=grid,
          values_grid=final_value_grid,
          discounting=discounting,
          boundary_condtions=[(None, upper_boundary_fn)],
      )[0]
  
      with self.subTest("OutputShape"):
          self.assertAllEqual(
              prices_estimated.shape.as_list(),
              (num_processes, num_strikes, num_grid_points),
          )
  
      # Compute prices at some locations for all strikes and processes
      loc_1 = 95
      loc_2 = 123
      loc_3 = 155
      inds = tf.stack([loc_1, loc_2, loc_3])
  
      gather_ind1 = tf.stack([[0, 1, 2], [0] * 3, inds], axis=-1)
      gather_ind2 = tf.stack([[0, 1, 2], [1] * 3, inds], axis=-1)
      gather_ind = tf.concat([gather_ind1, gather_ind2], axis=0)
  
      prices_at_locs = tf.gather_nd(prices_estimated, gather_ind)
      # Expected prices as computed with tff.black_scholes.option_price
      expected_prices = [25.726, 5.502, 6.3, 17.684, 11.221, 32.467]
      with self.subTest("Prices"):
          self.assertAllClose(
              prices_at_locs, expected_prices, rtol=1e-2, atol=1e-2
          )
  
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 6 for calling the GenericItoProcessTest obj's underlying member method code with example usage: obj.test_backward_pde_coeffs_with_batching_1d() and returns: unknown, unknown, tf, rates_expand