code: def _milstein_step(
      *,
      dim,
      i,
      written_count,
      current_state,
      result,
      drift_fn,
      volatility_fn,
      grad_volatility_fn,
      wiener_mean,
      num_samples,
      times,
      dt,
      sqrt_dt,
      keep_mask,
      random_type,
      seed,
      normal_draws,
      input_gradients,
      stratonovich_order,
      aux_normal_draws,
      record_samples,
  ):
      """Performs one step of Milstein scheme."""
      current_time = times[i + 1]
      written_count = tf.cast(written_count, tf.int32)
      if normal_draws is not None:
          dw = normal_draws[i]
      else:
          dw = random.mv_normal_sample(
              (num_samples,),
              mean=wiener_mean,
              random_type=random_type,
              seed=seed,
          )
      if aux_normal_draws is not None:
          stratonovich_draws = []
          for j in range(3):
              stratonovich_draws.append(
                  tf.reshape(
                      aux_normal_draws[j][i],
                      [num_samples, dim, stratonovich_order],
                  )
              )
      else:
          stratonovich_draws = []
          # Three sets of normal draws for stratonovich integrals.
          for j in range(3):
              stratonovich_draws.append(
                  random.mv_normal_sample(
                      (num_samples,),
                      mean=tf.zeros(
                          (dim, stratonovich_order),
                          dtype=current_state.dtype,
                          name="stratonovich_draws_{}".format(j),
                      ),
                      random_type=random_type,
                      seed=seed,
                  )
              )
  
      if dim == 1:
          drift = drift_fn(current_time, current_state)
          vol = volatility_fn(current_time, current_state)
          grad_vol = grad_volatility_fn(
              current_time, current_state, tf.ones_like(current_state)
          )
          next_state = _milstein_1d(
              dw=dw,
              dt=dt[i],
              sqrt_dt=sqrt_dt[i],
              current_state=current_state,
              drift=drift,
              vol=vol,
              grad_vol=grad_vol,
          )
      else:
          drift = drift_fn(current_time, current_state)
          vol = volatility_fn(current_time, current_state)
          # This is a list of size equal to the dimension of the state space `dim`.
          # It contains tensors of shape [num_samples, dim, wiener_dim] representing
          # the gradient of the volatility function. In our case, the dimension of the
          # wiener process `wiener_dim` is equal to the state dimension `dim`.
          grad_vol = [
              grad_volatility_fn(current_time, current_state, start)
              for start in input_gradients
          ]
          next_state = _milstein_nd(
              dim=dim,
              num_samples=num_samples,
              dw=dw,
              dt=dt[i],
              sqrt_dt=sqrt_dt[i],
              current_state=current_state,
              drift=drift,
              vol=vol,
              grad_vol=grad_vol,
              stratonovich_draws=stratonovich_draws,
              stratonovich_order=stratonovich_order,
          )
      if record_samples:
          result = result.write(written_count, next_state)
      else:
          result = next_state
      written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
  
      return i + 1, written_count, next_state, result
  
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 11 for calling a global function with example usage: _milstein_step() and returns: unknown,  written_count,  next_state,  result