code: def _milstein_hot(
      dim,
      vol,
      grad_vol,
      dt,
      sqrt_dt,
      dw,
      stratonovich_draws,
      stratonovich_order,
  ):
      """Higher order terms for Milstein update."""
      # Generate approximate Stratonovich integrals J(i,j) then replace the diagonal
      # with exact values.
      offdiag = _stratonovich_integral(
          dim=dim,
          dt=dt,
          sqrt_dt=sqrt_dt,
          dw=dw,
          stratonovich_draws=stratonovich_draws,
          order=stratonovich_order,
      )
      stratonovich_integrals = tf.linalg.set_diag(offdiag, dw * dw / 2)
  
      # Compute L_bar^{j1} b^{k, j2} J(j1, j2)
      # See Eq 3.4 of section 10.3 in [2]
      stacked_grad_vol = []
      for state_ix in range(dim):
          stacked_grad_vol.append(
              tf.transpose(
                  tf.stack([x[..., state_ix, :] for x in grad_vol], -1),
                  [0, 2, 1],
              )
          )
      stacked_grad_vol = tf.stack(stacked_grad_vol, 0)
      lbar = tf.matmul(stacked_grad_vol, vol)
      return tf.transpose(
          tf.reduce_sum(
              tf.multiply(lbar, stratonovich_integrals), [-2, -1]
          )
      )
  
imports: import functools
  import math
  from typing import Callable, List, Optional
  import numpy
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math import custom_loops
  from tf_quant_finance.math import gradient
  from tf_quant_finance.math import random_ops
  from tf_quant_finance.models import utils
label_desc: Use image node_id 7 for calling a global function with example usage: _milstein_hot(dim, vol, grad_vol, dt, sqrt_dt, dw, stratonovich_draws, stratonovich_order) and returns: tf