code: def _european_swaption_fd(
      batch_shape,
      model,
      exercise_times,
      fixed_leg_payment_times,
      fixed_leg_daycount_fractions,
      fixed_leg_coupon,
      notional,
      is_payer_swaption,
      time_step_fd,
      num_time_steps_fd,
      num_grid_points_fd,
      name,
      dtype,
  ):
      """Price European swaptions using finite difference."""
      with tf.name_scope(name):
          dim = model.dim()
          x_min = -0.5
          x_max = 0.5
          # grid.shape = (num_grid_points,2)
          grid = pde.grids.uniform_grid(
              minimums=[x_min] * dim,
              maximums=[x_max] * dim,
              sizes=[num_grid_points_fd] * dim,
              dtype=dtype,
          )
  
          # TODO(b/186876306): Remove dynamic shapes.
          pde_time_grid, pde_time_grid_dt = _create_pde_time_grid(
              exercise_times, time_step_fd, num_time_steps_fd, dtype
          )
          (
              maturities,
              unique_maturities,
              maturities_shape,
          ) = _create_term_structure_maturities(fixed_leg_payment_times)
  
          num_maturities = tf.shape(unique_maturities)[-1]
          x_meshgrid = _coord_grid_to_mesh_grid(grid)
          meshgrid_shape = tf.shape(x_meshgrid)
          broadcasted_maturities = tf.expand_dims(
              unique_maturities, axis=0
          )
  
          num_grid_points = tf.math.reduce_prod(meshgrid_shape[1:])
          shape_to_broadcast = tf.concat(
              [meshgrid_shape, [num_maturities]], axis=0
          )
  
          # Reshape `state_x`, `maturities` to (num_grid_points, num_maturities)
          state_x = tf.expand_dims(x_meshgrid, axis=-1)
          state_x = tf.broadcast_to(state_x, shape_to_broadcast)
          broadcasted_maturities = tf.broadcast_to(
              broadcasted_maturities, shape_to_broadcast[1:]
          )
  
          def _get_swap_payoff(payoff_time):
              broadcasted_exercise_times = tf.broadcast_to(
                  payoff_time, shape_to_broadcast[1:]
              )
  
              # Zero-coupon bond curve
              zcb_curve = model.discount_bond_price(
                  tf.transpose(
                      tf.reshape(
                          state_x,
                          [dim, num_grid_points * num_maturities],
                      )
                  ),
                  tf.reshape(broadcasted_exercise_times, [-1]),
                  tf.reshape(broadcasted_maturities, [-1]),
              )
              zcb_curve = tf.reshape(
                  zcb_curve, [num_grid_points, num_maturities]
              )
  
              maturities_index = tf.searchsorted(
                  unique_maturities, tf.reshape(maturities, [-1])
              )
  
              zcb_curve = tf.gather(
                  zcb_curve, maturities_index, axis=-1
              )
              # zcb_curve.shape = [num_grid_points] + [maturities_shape]
              zcb_curve = tf.reshape(
                  zcb_curve,
                  tf.concat(
                      [[num_grid_points], maturities_shape], axis=0
                  ),
              )
  
              # Shape after reduce_sum =
              # (num_grid_points, batch_shape)
              fixed_leg = tf.math.reduce_sum(
                  fixed_leg_coupon
                  * fixed_leg_daycount_fractions
                  * zcb_curve,
                  axis=-1,
              )
              float_leg = 1.0 - zcb_curve[..., -1]
              payoff_swap = float_leg - fixed_leg
              payoff_swap = tf.where(
                  is_payer_swaption, payoff_swap, -payoff_swap
              )
              return tf.reshape(
                  tf.transpose(payoff_swap),
                  tf.concat([batch_shape, meshgrid_shape[1:]], axis=0),
              )
  
          def _get_index(t, tensor_to_search):
              t = tf.expand_dims(t, axis=-1)
              index = tf.searchsorted(
                  tensor_to_search, t - _PDE_TIME_GRID_TOL, "right"
              )
              y = tf.gather(tensor_to_search, index)
              return tf.where(
                  tf.math.abs(t - y) < _PDE_TIME_GRID_TOL, index, -1
              )[0]
  
          sum_x_meshgrid = tf.math.reduce_sum(x_meshgrid, axis=0)
  
          def _is_exercise_time(t):
              return tf.reduce_any(
                  tf.where(
                      tf.math.abs(exercise_times[..., -1] - t)
                      < _PDE_TIME_GRID_TOL,
                      True,
                      False,
                  ),
                  axis=-1,
              )
  
          def _discounting_fn(t, grid):
              del grid
              f_0_t = model._instant_forward_rate_fn(
                  t
              )  # pylint: disable=protected-access
              return sum_x_meshgrid + f_0_t
  
          def _final_value():
              t = pde_time_grid[-1]
              payoff_swap = tf.nn.relu(_get_swap_payoff(t))
              is_ex_time = _is_exercise_time(t)
              return tf.where(
                  tf.reshape(
                      is_ex_time,
                      tf.concat([batch_shape, [1] * dim], axis=0),
                  ),
                  payoff_swap,
                  0.0,
              )
  
          def _values_transform_fn(t, grid, value_grid):
              zero = tf.zeros_like(value_grid)
              is_ex_time = _is_exercise_time(t)
  
              def _at_least_one_swaption_pays():
                  payoff_swap = tf.nn.relu(_get_swap_payoff(t))
                  return tf.where(
                      tf.reshape(
                          is_ex_time,
                          tf.concat([batch_shape, [1] * dim], axis=0),
                      ),
                      payoff_swap,
                      zero,
                  )
  
              v_star = tf.cond(
                  tf.math.reduce_any(is_ex_time),
                  _at_least_one_swaption_pays,
                  lambda: zero,
              )
              return grid, tf.maximum(value_grid, v_star)
  
          def _pde_time_step(t):
              index = _get_index(t, pde_time_grid)
              return pde_time_grid_dt[index]
  
          # Use default boundary conditions, d^2V/dx_i^2 = 0
          boundary_conditions = [(None, None) for i in range(dim)]
          # res[0] contains the swaption prices.
          # res[0].shape = batch_shape + [num_grid_points] * dim
          res = model.fd_solver_backward(
              pde_time_grid[-1],
              0.0,
              grid,
              values_grid=_final_value(),
              time_step=_pde_time_step,
              boundary_conditions=boundary_conditions,
              values_transform_fn=_values_transform_fn,
              discounting=_discounting_fn,
              dtype=dtype,
          )
  
          idx = tf.searchsorted(
              tf.convert_to_tensor(grid),
              tf.expand_dims(
                  tf.convert_to_tensor([0.0] * dim, dtype=dtype),
                  axis=-1,
              ),
          )
          # idx.shape = (dim, 1)
          idx = (
              tf.squeeze(idx) if dim > 1 else tf.reshape(idx, shape=[1])
          )
          slices = [slice(None)] + [
              slice(i, i + 1) for i in tf.unstack(idx)
          ]
          # shape = batch_shape + [1] * dim
          option_value = res[0][slices]
          # shape = batch_shape
          option_value = tf.squeeze(
              option_value, axis=list(range(-dim, 0))
          )
  
          # output_shape = batch_shape
          return notional * tf.reshape(option_value, batch_shape)
  
file_path: tf-quant-finance/tf_quant_finance/models/hjm/swaption_pricing.py
imports: from typing import Callable, Union
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance.math import pde
  from tf_quant_finance.math import random
  from tf_quant_finance.models import utils
  from tf_quant_finance.models import valuation_method
  from tf_quant_finance.models.hjm import gaussian_hjm
  from tf_quant_finance.models.hjm import quasi_gaussian_hjm
  from tf_quant_finance.models.hjm import swaption_util
label_desc: Use image node_id 3 for calling a global function with example usage: _european_swaption_fd(batch_shape, model, exercise_times, fixed_leg_payment_times, fixed_leg_daycount_fractions, fixed_leg_coupon, notional, is_payer_swaption, time_step_fd, num_time_steps_fd, num_grid_points_fd, name, dtype) and returns: unknown, tf, tf, tf, unknown, tf, grid,  tf, pde_time_grid_dt, tf