code: def _prepare_swaption_indices(tensor_shape):
      """Indices for `gather_nd` for analytic valuation.
  
      For a `Tensor` x of shape `tensor_shape` = [n] + batch_shape + [n], this
      function returns indices for tf.gather_nd to get `x[i,...,i]`
  
      Args:
        tensor_shape: A list of length `k` representing shape of the `Tensor`.
  
      Returns:
        A `Tensor` of shape (num_elements, k) where num_elements= n * batch_size
        of dtype tf.int64.
      """
  
      tensor_shape = np.array(tensor_shape, dtype=np.int64)
      batch_shape = tensor_shape[1:-1]
      batch_size = np.prod(batch_shape)
      index_list = []
      for i in range(len(tensor_shape)):
          index = np.arange(0, tensor_shape[i], dtype=np.int64)
          if i == 0 or i == len(tensor_shape) - 1:
              index = tf.tile(index, [batch_size])
          else:
              index = np.tile(
                  np.repeat(index, np.prod(tensor_shape[i + 1 :])),
                  [np.prod(tensor_shape[1:i])],
              )
          index_list.append(index)
  
      return tf.stack(index_list, axis=-1)
  
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 4 for calling a global function with example usage: _prepare_swaption_indices(tensor_shape) and returns: tf