code: def _leverage_function_using_pde(
      *,
      risk_free_rate,
      dividend_yield,
      lv_model,
      variance_model,
      rho,
      initial_spot,
      initial_variance,
      max_time,
      time_step,
      num_grid_points,
      grid_minimums,
      grid_maximums,
      dtype,
  ):
      """Computes Leverage function using Fokker-Planck PDE for joint density.
  
      This function computes the leverage function for the LSV model by first
      computing the joint probablity density function `p(t, X(t), v(t))` where
      `X(t)` is the log of the spot price and `v(t)` is the variance at time `t`.
      The joint probablity density is computed using the Fokker-Planck equation of
      the LSV model (see 6.8.2 in Ref [1]):
      ```None
      dp/dt = 1/2 d^2 [v L(t,X)^2 p]/dX^2 + 1/2 d^2 [b(v)^2 p]/dv^2 +
              rho d^2 [sqrt(v)L(t,X)b(v) p]/dXdv - d[(r - d - 1/2 v L(t,X)^2)p]/dX -
              d[a(v) p]/dv
      ```
  
      where `a(v)` and `b(v)` are the drift and diffusion functions for the
      variance process. Defining
  
      ```None
      I_n(k,t) = int v^n p(t, k, v) dv
      ```
  
      we can calculate the leverage function as follows:
      ```None
      L(k, t) = sigma(exp(k), t) sqrt(I_0(k, t)/I_1(k, t)).
      ```
  
      Args:
        risk_free_rate: A scalar real `Tensor` specifying the (continuosly
          compounded) risk free interest rate. If the underlying is an FX rate, then
          use this input to specify the domestic interest rate.
        dividend_yield: A real scalar `Tensor` specifying the (continuosly
          compounded) dividend yield. If the underlying is an FX rate, then use this
          input to specify the foreign interest rate.
        lv_model: An instance of `LocalVolatilityModel` specifying the local
          volatility for the spot price.
        variance_model: An instance of `LSVVarianceModel` specifying the dynamics of
          the variance process of the LSV model.
        rho: A real scalar `Tensor` specifying the correlation between spot price
          and the stochastic variance.
        initial_spot: A real scalar `Tensor` specifying the underlying spot price on
          the valuation date.
        initial_variance: A real scalar `Tensor` specifying the initial variance on
          the valuation date.
        max_time: A real scalar `Tensor` specifying the maximum time to which the
          Fokker-Planck PDE is evolved.
        time_step: A real scalar `Tensor` specifying the time step during the
          numerical solution of the Fokker-Planck PDE.
        num_grid_points: A scalar integer `Tensor` specifying the number of
          discretization points for each spatial dimension.
        grid_minimums: An optional `Tensor` of size 2 containing the minimum grid
          points for PDE spatial discretization. `grid_minimums[0]` correspond
          to the minimum spot price in the spatial grid and `grid_minimums[1]`
          correspond to the minimum variance value.
        grid_maximums: An optional `Tensor` of size 2 containing the maximum grid
          points for PDE spatial discretization. `grid_maximums[0]` correspond
          to the maximum spot price in the spatial grid and `grid_maximums[1]`
          correspond to the maximum variance value.
        dtype: The default dtype to use when converting values to `Tensor`s.
  
      Returns:
        A Python callable which computes the Leverage function `L(t, S(t))`. The
        function accepts a scalar `Tensor` corresponding to time 't' and a real
        `Tensor` of shape `[num_samples, 1]` corresponding to the spot price (S) as
        inputs  and return a real `Tensor` corresponding to the leverage function
        computed at (S,t).
  
      """
      if variance_model.dim() > 1:
          raise ValueError(
              "The default model of Leverage function doesn't support "
              "the variance process with more than 1 factor."
          )
  
      pde_grid_tol = _machine_eps(dtype)
      rho = tf.convert_to_tensor(rho, dtype=dtype)
      initial_spot = tf.convert_to_tensor(initial_spot, dtype=dtype)
      initial_log_spot = tf.math.log(
          tf.convert_to_tensor(initial_spot, dtype=dtype)
      )
      initial_variance = tf.convert_to_tensor(
          initial_variance, dtype=dtype
      )
      risk_free_rate = tf.convert_to_tensor(risk_free_rate, dtype=dtype)
      dividend_yield = tf.convert_to_tensor(dividend_yield, dtype=dtype)
      rho = tf.convert_to_tensor(rho, dtype=dtype)
  
      x_scale = initial_log_spot
      y_scale = initial_variance
      # scaled log spot = log(spot/initial_spot)
      # scaled variance = variance / initial_variance
      scaled_initial_point = tf.convert_to_tensor(
          [0.0, 1.0], dtype=dtype
      )
  
      # These are minimums and maximums for scaled log spot and scaled variance
      if grid_minimums is None:
          grid_minimums = [0.01, 0.0001]
      else:
          grid_minimums = tf.convert_to_tensor(
              grid_minimums, dtype=dtype
          )
          grid_minimums = [
              grid_minimums[0] / initial_spot,
              grid_minimums[1] / initial_variance,
          ]
      if grid_maximums is None:
          grid_maximums = [10.0, 5.0]
      else:
          grid_maximums = tf.convert_to_tensor(
              grid_maximums, dtype=dtype
          )
          grid_maximums = [
              grid_maximums[0] / initial_spot,
              grid_maximums[1] / initial_variance,
          ]
  
      log_spot_min = tf.math.log(
          tf.convert_to_tensor([grid_minimums[0]], dtype=dtype)
      )
      log_spot_max = tf.math.log(
          tf.convert_to_tensor([grid_maximums[0]], dtype=dtype)
      )
      variance_min = tf.convert_to_tensor(
          [grid_minimums[1]], dtype=dtype
      )
      variance_max = tf.convert_to_tensor(
          [grid_maximums[1]], dtype=dtype
      )
  
      grid_minimums = tf.concat([log_spot_min, variance_min], axis=0)
      grid_maximums = tf.concat([log_spot_max, variance_max], axis=0)
  
      grid = _tavella_randell_nonuniform_grid(
          grid_minimums,
          grid_maximums,
          scaled_initial_point,
          num_grid_points,
          0.3,
          dtype,
      )
      grid = [
          tf.expand_dims(grid[0], axis=0),
          tf.expand_dims(grid[1], axis=0),
      ]
  
      delta_x = tf.math.reduce_min(grid[0][0, 1:] - grid[0][0, :-1])
      delta_y = tf.math.reduce_min(grid[1][0, 1:] - grid[1][0, :-1])
      # Initialize leverage function L(t=0, S) = 1
      leverage_fn = functools.partial(
          linear.interpolate,
          x_data=[[0.0, 1.0]],
          y_data=[[1.0, 1.0]],
          dtype=dtype,
      )
  
      def _initial_value():
          """Computes initial value as a delta function delta(log_spot(t), var(0))."""
          log_spot, variance = tf.meshgrid(*grid)
  
          init_value = tf.where(
              tf.math.logical_and(
                  tf.math.abs(log_spot - scaled_initial_point[0])
                  < delta_x + pde_grid_tol,
                  tf.math.abs(variance - scaled_initial_point[1])
                  < delta_y + pde_grid_tol,
              ),
              1.0 / (delta_x * delta_y * 4),
              0.0,
          )
          # initial_value.shape = (1, num_grid_x, num_grid_y)
          return tf.expand_dims(init_value, axis=0)
  
      def _second_order_coeff_fn(t, grid):
          log_spot = grid[0] + x_scale
          variance = grid[1] * y_scale
          leverage_fn_t_x = leverage_fn(log_spot)
          val_xx = 0.5 * variance * leverage_fn_t_x**2
          val_xy = (
              0.5
              * (
                  rho
                  * tf.math.sqrt(variance)
                  * leverage_fn_t_x
                  * variance_model.volatility_fn()(t, variance)
              )
              / y_scale
          )
          val_yx = val_xy
          val_yy = (
              0.5
              * variance_model.volatility_fn()(t, variance) ** 2
              / y_scale**2
          )
          # return list of shape = (2,2). Each element has shape = grid.shape
          return [[-val_yy, -val_yx], [-val_xy, -val_xx]]
  
      def _first_order_coeff_fn(t, grid):
          log_spot = grid[0] + x_scale
          variance = grid[1] * y_scale
          leverage_fn_t_x = leverage_fn(log_spot)
          val_x = (
              risk_free_rate
              - dividend_yield
              - 0.5 * variance * leverage_fn_t_x**2
          )
          val_y = variance_model.drift_fn()(t, variance)
          # return list of shape = (2,). Each element has shape = grid.shape
          return [val_y / y_scale, val_x]
  
      def _compute_leverage_fn(t, coord_grid, value_grid):
          log_spot = tf.expand_dims(coord_grid[0], axis=-1) + x_scale
          local_volatility_values = lv_model.local_volatility_fn()(
              t, tf.math.exp(log_spot)
          )
          # TODO(b/176826650): Large values represent instability. Eventually this
          # should be addressed inside local vol model.
          local_volatility_values = tf.where(
              tf.math.abs(local_volatility_values) > 1e4,
              tf.constant(0.0, dtype=dtype),
              local_volatility_values,
          )
          # variance_given_logspot.shape = (num_grid_x, 1)
          variance_given_logspot = (
              _conditional_expected_variance_from_pde_solution(
                  [coord_grid[0] + x_scale, coord_grid[1] * y_scale],
                  value_grid,
                  dtype,
              )(log_spot)
          )
  
          leverage_fn_values = tf.math.divide_no_nan(
              local_volatility_values,
              tf.math.sqrt(variance_given_logspot),
          )
  
          leverage_fn = functools.partial(
              linear.interpolate,
              x_data=grid[0] + x_scale,
              y_data=tf.transpose(leverage_fn_values),
              dtype=dtype,
          )
  
          return leverage_fn
  
      @pde.boundary_conditions.neumann
      def _trivial_neumann_boundary(t, location_grid):
          del t, location_grid
          return 0.0
  
      # joint_density.shape = (1, num_grid_x, num_grid_y)
      joint_density = _initial_value()
  
      def loop_body(i, tstart, joint_density, leverage_fn_values):
          (
              joint_density,
              coord_grid,
              _,
              _,
          ) = pde.fd_solvers.solve_forward(
              tstart,
              tstart + time_step,
              coord_grid=[grid[0][0], grid[1][0]],
              values_grid=joint_density,
              time_step=time_step / 10.0,
              values_transform_fn=None,
              inner_second_order_coeff_fn=_second_order_coeff_fn,
              inner_first_order_coeff_fn=_first_order_coeff_fn,
              zeroth_order_coeff_fn=None,
              boundary_conditions=[
                  [
                      _trivial_neumann_boundary,
                      _trivial_neumann_boundary,
                  ],
                  [
                      _trivial_neumann_boundary,
                      _trivial_neumann_boundary,
                  ],
              ],
              dtype=dtype,
          )
          joint_density = tf.math.maximum(joint_density, 0.0)
          area_under_joint_density = _two_d_integration(
              [grid[0][0, :], grid[1][0, :]], joint_density
          )
          joint_density = joint_density / area_under_joint_density
  
          # TODO(b/176826743): Perform fixed point iteration instead of one step
          # update
          leverage_fn = _compute_leverage_fn(
              tstart + time_step, coord_grid, joint_density
          )
          leverage_v = leverage_fn(grid[0][0, :] + x_scale)[0, :]
          leverage_fn_values = leverage_fn_values.write(i, leverage_v)
  
          return (
              i + 1,
              tstart + time_step,
              joint_density,
              leverage_fn_values,
          )
  
      times = tf.range(
          0.0, max_time + time_step, time_step, dtype=dtype
      )
      tstart = times[0]
      first_leverage_value = leverage_fn(grid[0][0])[0]
      leverage_fn_values = tf.TensorArray(
          dtype=dtype,
          size=tff_utils.get_shape(times)[0],
          element_shape=tff_utils.get_shape(first_leverage_value),
          clear_after_read=False,
      )
      leverage_fn_values.write(0, first_leverage_value)
  
      loop_cond = lambda i, tstart, *args: tf.less(tstart, max_time)
      initial_args = (1, tstart, joint_density, leverage_fn_values)
      _, _, _, leverage_fn_values = tf.while_loop(
          loop_cond, loop_body, initial_args
      )
      leverage_fn_values = leverage_fn_values.stack()
  
      # leverage_fn_values.shape = (num_pde_timesteps, num_grid_x,)
      leverage_fn_values = tf.convert_to_tensor(
          leverage_fn_values, dtype=dtype
      )
  
      def _return_fn(t, spot):
          leverage_fn_interpolator = (
              math.interpolation.interpolation_2d.Interpolation2D(
                  x_data=[times],
                  y_data=tf.expand_dims(
                      tf.repeat(
                          grid[0] + x_scale, times.shape[0], axis=0
                      ),
                      axis=0,
                  ),
                  z_data=tf.expand_dims(leverage_fn_values, axis=0),
                  dtype=dtype,
              )
          )
          return leverage_fn_interpolator.interpolate(
              t, tf.math.log(spot)
          )
  
      return _return_fn
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tf_quant_finance import datetime
  from tf_quant_finance import math
  from tf_quant_finance import utils
  from tf_quant_finance.experimental import local_volatility
  from tf_quant_finance.experimental.pricing_platform.framework.market_data import utils
  from tf_quant_finance.math import pde
  from tf_quant_finance.math.interpolation import linear
  from tf_quant_finance.models import generic_ito_process
label_desc: Use image node_id 6 for calling a global function with example usage: _leverage_function_using_pde() and returns: _return_fn, tf, list, list, leverage_fn, int, unknown,  unknown,  joint_density,  leverage_fn_values, leverage_fn_interpolator