code: def _bermudan_swaption_fd(
      batch_shape,
      model,
      exercise_times,
      unique_exercise_times,
      fixed_leg_payment_times,
      fixed_leg_daycount_fractions,
      fixed_leg_coupon,
      notional,
      is_payer_swaption,
      time_step_fd,
      num_grid_points_fd,
      name,
      dtype,
  ):
      """Price Bermudan swaptions using finite difference."""
      with tf.name_scope(name):
          longest_exercise_time = unique_exercise_times[-1]
          if time_step_fd is None:
              time_step_fd = longest_exercise_time / 100.0
          short_rate_min = -0.2
          short_rate_max = 0.2
          # grid.shape=(num_grid_points,2)
          grid = pde.grids.uniform_grid(
              minimums=[short_rate_min],
              maximums=[short_rate_max],
              sizes=[num_grid_points_fd],
              dtype=dtype,
          )
  
          pde_time_grid = tf.concat(
              [
                  unique_exercise_times,
                  tf.range(
                      0.0,
                      longest_exercise_time,
                      time_step_fd,
                      dtype=dtype,
                  ),
              ],
              axis=0,
          )
          # This time grid is now sorted and contains the Bermudan exercise times
          pde_time_grid = tf.sort(
              pde_time_grid, name="sort_pde_time_grid"
          )
          pde_time_grid_dt = pde_time_grid[1:] - pde_time_grid[:-1]
          pde_time_grid_dt = tf.concat(
              [[100.0], pde_time_grid_dt], axis=-1
          )
          # Remove duplicates.
          mask = tf.math.greater(pde_time_grid_dt, _PDE_TIME_GRID_TOL)
          pde_time_grid = tf.boolean_mask(pde_time_grid, mask)
          pde_time_grid_dt = tf.boolean_mask(pde_time_grid_dt, mask)
  
          maturities = fixed_leg_payment_times
          maturities_shape = maturities.shape
  
          unique_maturities, _ = tf.unique(
              tf.reshape(maturities, shape=[-1])
          )
          unique_maturities = tf.sort(
              unique_maturities, name="sort_maturities"
          )
  
          num_exercise_times = tf.shape(pde_time_grid)[-1]
          num_maturities = tf.shape(unique_maturities)[-1]
  
          short_rates = tf.reshape(grid[0], grid[0].shape + [1, 1])
          broadcasted_exercise_times = tf.reshape(
              pde_time_grid, [1] + pde_time_grid.shape + [1]
          )
          broadcasted_maturities = tf.reshape(
              unique_maturities, [1, 1] + unique_maturities.shape
          )
  
          # Reshape `short_rate`, `exercise_times` and `maturities` to
          # (num_grid_points, num_exercise_times, num_maturities)
          short_rates = tf.broadcast_to(
              short_rates,
              grid[0].shape + [num_exercise_times, num_maturities],
          )
          broadcasted_exercise_times = tf.broadcast_to(
              broadcasted_exercise_times,
              grid[0].shape + [num_exercise_times, num_maturities],
          )
          broadcasted_maturities = tf.broadcast_to(
              broadcasted_maturities,
              grid[0].shape + [num_exercise_times, num_maturities],
          )
  
          # Zero-coupon bond curve
          zcb_curve = model.discount_bond_price(
              tf.expand_dims(short_rates, axis=-1),
              broadcasted_exercise_times,
              broadcasted_maturities,
          )[..., 0]
  
          exercise_times_index = tf.searchsorted(
              pde_time_grid, tf.reshape(exercise_times, [-1])
          )
          maturities_index = tf.searchsorted(
              unique_maturities, tf.reshape(maturities, [-1])
          )
  
          # gather_index.shape = (num_grid_points * np.cumprod(maturities_shape), 3)
          gather_index = _prepare_indices_ijj(
              tf.range(0, num_grid_points_fd),
              exercise_times_index,
              maturities_index,
          )
          zcb_curve = tf.gather_nd(zcb_curve, gather_index)
          # zcb_curve.shape = [num_grid_points_fd] + [maturities_shape]
          zcb_curve = tf.reshape(
              zcb_curve, [num_grid_points_fd] + maturities_shape
          )
          # Shape after reduce_sum=(num_grid_points, batch_shape, num_exercise_times)
          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_at_exercise = float_leg - fixed_leg
          payoff_at_exercise = tf.where(
              is_payer_swaption, payoff_at_exercise, -payoff_at_exercise
          )
  
          unrepeated_exercise_times = exercise_times[..., -1]
          exercise_times_index = tf.searchsorted(
              pde_time_grid, tf.reshape(unrepeated_exercise_times, [-1])
          )
          _, payoff_swap = _map_payoff_to_sim_times(
              tf.reshape(
                  exercise_times_index, unrepeated_exercise_times.shape
              ),
              payoff_at_exercise,
              num_grid_points_fd,
          )
  
          # payoff_swap.shape=(num_grid_points_fd, batch_shape, num_exercise_times).
          # Transpose so that the [num_grid_points_fd] is the last dimension; this is
          # needed for broadcasting inside the PDE solver.
          payoff_swap = tf.transpose(payoff_swap)
  
          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]
  
          def _second_order_coeff_fn(t, grid):
              del grid
              return [[model.volatility(t) ** 2 / 2]]
  
          def _first_order_coeff_fn(t, grid):
              s = grid[0]
              return [model.drift_fn()(t, s)]
  
          def _zeroth_order_coeff_fn(t, grid):
              del t
              return -grid[0]
  
          @pde.boundary_conditions.dirichlet
          def _lower_boundary_fn(t, grid):
              del grid
              index = _get_index(t, pde_time_grid)
              result = tf.where(
                  index > -1, payoff_swap[index, ..., 0], 0.0
              )
              return tf.where(is_payer_swaption, 0.0, result)
  
          @pde.boundary_conditions.dirichlet
          def _upper_boundary_fn(t, grid):
              del grid
              index = _get_index(t, pde_time_grid)
              result = tf.where(
                  index > -1, payoff_swap[index, ..., 0], 0.0
              )
              return tf.where(is_payer_swaption, result, 0.0)
  
          def _final_value():
              return tf.nn.relu(payoff_swap[-1])
  
          def _values_transform_fn(t, grid, value_grid):
              index = _get_index(t, pde_time_grid)
              v_star = tf.where(
                  index > -1, tf.nn.relu(payoff_swap[index]), 0.0
              )
              return grid, tf.maximum(value_grid, v_star)
  
          def _pde_time_step(t):
              index = _get_index(t, pde_time_grid)
              dt = pde_time_grid_dt[index]
              return dt
  
          res = pde.fd_solvers.solve_backward(
              longest_exercise_time,
              0.0,
              grid,
              values_grid=_final_value(),
              time_step=_pde_time_step,
              boundary_conditions=[
                  (_lower_boundary_fn, _upper_boundary_fn)
              ],
              values_transform_fn=_values_transform_fn,
              second_order_coeff_fn=_second_order_coeff_fn,
              first_order_coeff_fn=_first_order_coeff_fn,
              zeroth_order_coeff_fn=_zeroth_order_coeff_fn,
              dtype=dtype,
          )
  
          r0 = model.instant_forward_rate(0.0)
          option_value = linear.interpolate(r0, res[1], res[0])
          return tf.reshape(
              notional * tf.transpose(option_value), batch_shape
          )
  
imports: from typing import Callable, Union
  import numpy
  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.math.interpolation import linear
  from tf_quant_finance.math.root_search import brent
  from tf_quant_finance.models import utils
  from tf_quant_finance.models.hjm import swaption_util
  from tf_quant_finance.models.hull_white import one_factor
  from tf_quant_finance.models.hull_white import zero_coupon_bond_option
  from tf_quant_finance.models.longstaff_schwartz import lsm
label_desc: Use image node_id 9 for calling a global function with example usage: _bermudan_swaption_fd(batch_shape, model, exercise_times, unique_exercise_times, fixed_leg_payment_times, fixed_leg_daycount_fractions, fixed_leg_coupon, notional, is_payer_swaption, time_step_fd, num_grid_points_fd, name, dtype) and returns: tf, tf, list, list, unknown, tf, tf, tf, grid,  tf, dt