code: def _sample(
      *,
      dim,
      drift_fn,
      volatility_fn,
      grad_volatility_fn,
      times,
      time_step,
      keep_mask,
      num_requested_times,
      num_samples,
      initial_state,
      random_type,
      seed,
      swap_memory,
      skip,
      precompute_normal_draws,
      watch_params,
      time_indices,
      input_gradients,
      stratonovich_order,
      dtype,
  ):
      """Returns a sample of paths from the process using the Milstein method."""
      dt = times[1:] - times[:-1]
      sqrt_dt = tf.sqrt(dt)
      current_state = initial_state + tf.zeros(
          [num_samples, dim], dtype=initial_state.dtype
      )
      if dt.shape.is_fully_defined():
          steps_num = dt.shape.as_list()[-1]
      else:
          steps_num = tf.shape(dt)[-1]
      # In order to use low-discrepancy random_type we need to generate the sequence
      # of independent random normals upfront. We also precompute random numbers
      # for stateless random type in order to ensure independent samples for
      # multiple function calls with different seeds.
      if precompute_normal_draws or random_type in (
          random.RandomType.SOBOL,
          random.RandomType.HALTON,
          random.RandomType.HALTON_RANDOMIZED,
          random.RandomType.STATELESS,
          random.RandomType.STATELESS_ANTITHETIC,
      ):
          # Process dimension plus auxiliary random variables for stratonovich
          # integral computation.
          all_normal_draws = utils.generate_mc_normal_draws(
              num_normal_draws=dim + 3 * dim * stratonovich_order,
              num_time_steps=steps_num,
              num_sample_paths=num_samples,
              random_type=random_type,
              dtype=dtype,
              seed=seed,
              skip=skip,
          )
          normal_draws = all_normal_draws[:, :, :dim]
          wiener_mean = None
          # Auxiliary normal draws for use with the stratonovich integral
          # approximation.
          aux_normal_draws = []
          start = dim
          for _ in range(3):
              end = start + dim * stratonovich_order
              aux_normal_draws.append(all_normal_draws[:, :, start:end])
              start = end
      else:
          # If pseudo or anthithetic sampling is used, proceed with random sampling
          # at each step.
          wiener_mean = tf.zeros(
              (dim,), dtype=dtype, name="wiener_mean"
          )
          normal_draws = None
          aux_normal_draws = None
      if watch_params is None:
          # Use while_loop if `watch_params` is not passed
          return _while_loop(
              dim=dim,
              steps_num=steps_num,
              current_state=current_state,
              drift_fn=drift_fn,
              volatility_fn=volatility_fn,
              grad_volatility_fn=grad_volatility_fn,
              wiener_mean=wiener_mean,
              num_samples=num_samples,
              times=times,
              dt=dt,
              sqrt_dt=sqrt_dt,
              time_step=time_step,
              keep_mask=keep_mask,
              num_requested_times=num_requested_times,
              swap_memory=swap_memory,
              random_type=random_type,
              seed=seed,
              normal_draws=normal_draws,
              input_gradients=input_gradients,
              stratonovich_order=stratonovich_order,
              aux_normal_draws=aux_normal_draws,
              dtype=dtype,
          )
      else:
          # Use custom for_loop if `watch_params` is specified
          return _for_loop(
              dim=dim,
              steps_num=steps_num,
              current_state=current_state,
              drift_fn=drift_fn,
              volatility_fn=volatility_fn,
              grad_volatility_fn=grad_volatility_fn,
              wiener_mean=wiener_mean,
              num_samples=num_samples,
              times=times,
              dt=dt,
              sqrt_dt=sqrt_dt,
              time_indices=time_indices,
              keep_mask=keep_mask,
              watch_params=watch_params,
              random_type=random_type,
              seed=seed,
              normal_draws=normal_draws,
              input_gradients=input_gradients,
              stratonovich_order=stratonovich_order,
              aux_normal_draws=aux_normal_draws,
          )
  
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 2 for calling a global function with example usage: _sample() and returns: _while_loop, _for_loop