code: def sample(
      *,
      dim: int,
      drift_fn: Callable[..., types.RealTensor],
      volatility_fn: Callable[..., types.RealTensor],
      times: types.RealTensor,
      time_step: Optional[types.RealTensor] = None,
      num_time_steps: Optional[types.IntTensor] = None,
      num_samples: types.IntTensor = 1,
      initial_state: Optional[types.RealTensor] = None,
      grad_volatility_fn: Optional[
          Callable[..., List[types.RealTensor]]
      ] = None,
      random_type: Optional[random.RandomType] = None,
      seed: Optional[types.IntTensor] = None,
      swap_memory: bool = True,
      skip: types.IntTensor = 0,
      precompute_normal_draws: bool = True,
      watch_params: Optional[List[types.RealTensor]] = None,
      stratonovich_order: int = 5,
      dtype: Optional[tf.DType] = None,
      name: Optional[str] = None,
  ) -> types.RealTensor:
      r"""Returns a sample paths from the process using the Milstein method.
  
      For an Ito process,
  
      ```
        dX = a(t, X_t) dt + b(t, X_t) dW_t
      ```
      given drift `a`, volatility `b` and derivative of volatility `b'`, the
      Milstein method generates a
      sequence {Y_n} approximating X
  
      ```
      Y_{n+1} = Y_n + a(t_n, Y_n) dt + b(t_n, Y_n) dW_n + \frac{1}{2} b(t_n, Y_n)
      b'(t_n, Y_n) ((dW_n)^2 - dt)
      ```
      where `dt = t_{n+1} - t_n`, `dW_n = (N(0, t_{n+1}) - N(0, t_n))` and `N` is a
      sample from the Normal distribution.
  
      In higher dimensions, when `a(t, X_t)` is a d-dimensional vector valued
      function and `W_t` is a d-dimensional Wiener process, we have for the kth
      element of the expansion:
  
      ```
      Y_{n+1}[k] = Y_n[k] + a(t_n, Y_n)[k] dt + \sum_{j=1}^d b(t_n, Y_n)[k, j]
      dW_n[j] + \sum_{j_1=1}^d \sum_{j_2=1}^d L_{j_1} b(t_n, Y_n)[k, j_2] I(j_1,
      j_2)
      ```
      where `L_{j} = \sum_{i=1}^d b(t_n, Y_n)[i, j] \frac{\partial}{\partial x^i}`
      is an operator and `I(j_1, j_2) = \int_{t_n}^{t_{n+1}} \int_{t_n}^{s_1}
      dW_{s_2}[j_1] dW_{s_1}[j_2]` is a multiple Ito integral.
  
  
      See [1] and [2] for details.
  
      #### References
      [1]: Wikipedia. Milstein method:
      https://en.wikipedia.org/wiki/Milstein_method
      [2]: Peter E. Kloeden,  Eckhard Platen. Numerical Solution of Stochastic
        Differential Equations. Springer. 1992
  
      Args:
        dim: Python int greater than or equal to 1. The dimension of the Ito
          Process.
        drift_fn: A Python callable to compute the drift of the process. The
          callable should accept two real `Tensor` arguments of the same dtype. The
          first argument is the scalar time t, the second argument is the value of
          Ito process X - tensor of shape `batch_shape + [dim]`. The result is
          value of drift a(t, X). The return value of the callable is a real
          `Tensor` of the same dtype as the input arguments and of shape
          `batch_shape + [dim]`.
        volatility_fn: A Python callable to compute the volatility of the process.
          The callable should accept two real `Tensor` arguments of the same dtype
          as `times`. The first argument is the scalar time t, the second argument
          is the value of Ito process X - tensor of shape `batch_shape + [dim]`. The
          result is value of volatility b(t, X). The return value of the callable is
          a real `Tensor` of the same dtype as the input arguments and of shape
          `batch_shape + [dim, dim]`.
        times: Rank 1 `Tensor` of increasing positive real values. The times at
          which the path points are to be evaluated.
        time_step: An optional scalar real `Tensor` - maximal distance between
          points in grid in Milstein schema.
          Either this or `num_time_steps` should be supplied.
          Default value: `None`.
        num_time_steps: An optional Scalar integer `Tensor` - a total number of time
          steps performed by the algorithm. The maximal distance between points in
          grid is bounded by `times[-1] / (num_time_steps - times.shape[0])`.
          Either this or `time_step` should be supplied.
          Default value: `None`.
        num_samples: Positive scalar `int`. The number of paths to draw.
          Default value: 1.
        initial_state: `Tensor` of shape `[dim]`. The initial state of the
          process.
          Default value: None which maps to a zero initial state.
        grad_volatility_fn: An optional python callable to compute the gradient of
          `volatility_fn`. The callable should accept three real `Tensor` arguments
          of the same dtype as `times`. The first argument is the scalar time t. The
          second argument is the value of Ito process X - tensor of shape
          `batch_shape + [dim]`. The third argument is a tensor of input gradients
          of shape `batch_shape + [dim]` to pass to `gradient.fwd_gradient`. The
          result is a list of values corresponding to the forward gradient of
          volatility b(t, X) with respect to X. The return value of the callable is
          a list of size `dim` containing real `Tensor`s of the same dtype as the
          input arguments and of shape `batch_shape + [dim, dim]`. Each index of the
          list corresponds to a dimension of the state. If `None`, the gradient is
          computed from `volatility_fn` using forward differentiation.
        random_type: Enum value of `RandomType`. The type of (quasi)-random number
          generator to use to generate the paths.
          Default value: None which maps to the standard pseudo-random numbers.
        seed: Seed for the random number generator. The seed is only relevant if
          `random_type` is one of `[STATELESS, PSEUDO, HALTON_RANDOMIZED,
          PSEUDO_ANTITHETIC, STATELESS_ANTITHETIC]`. For `PSEUDO`,
          `PSEUDO_ANTITHETIC` and `HALTON_RANDOMIZED` the seed should be a Python
          integer. For `STATELESS` and  `STATELESS_ANTITHETIC `must be supplied as
          an integer `Tensor` of shape `[2]`.
          Default value: `None` which means no seed is set.
        swap_memory: A Python bool. Whether GPU-CPU memory swap is enabled for this
          op. See an equivalent flag in `tf.while_loop` documentation for more
          details. Useful when computing a gradient of the op since `tf.while_loop`
          is used to propagate stochastic process in time.
          Default value: True.
        skip: `int32` 0-d `Tensor`. The number of initial points of the Sobol or
          Halton sequence to skip. Used only when `random_type` is 'SOBOL',
          'HALTON', or 'HALTON_RANDOMIZED', otherwise ignored.
          Default value: `0`.
        precompute_normal_draws: Python bool. Indicates whether the noise increments
          `N(0, t_{n+1}) - N(0, t_n)` are precomputed. For `HALTON` and `SOBOL`
          random types the increments are always precomputed. While the resulting
          graph consumes more memory, the performance gains might be significant.
          Default value: `True`.
        watch_params: An optional list of zero-dimensional `Tensor`s of the same
          `dtype` as `initial_state`. If provided, specifies `Tensor`s with respect
          to which the differentiation of the sampling function will happen. A more
          efficient algorithm is used when `watch_params` are specified. Note the
          the function becomes differentiable only wrt to these `Tensor`s and the
          `initial_state`. The gradient wrt any other `Tensor` is set to be zero.
        stratonovich_order: A positive integer. The number of terms to use when
          calculating the approximate Stratonovich integrals in the multidimensional
          scheme. Stratonovich integrals are an alternative to Ito integrals, and
          can be used interchangeably when defining the higher order terms in the
          update equation. We use Stratonovich integrals here because they have a
          convenient approximation scheme for calculating cross terms involving
          different components of the Wiener process. See Eq. 8.10 in Section 5.8 of
          [2]. Default value: `5`.
        dtype: `tf.Dtype`. If supplied the dtype for the input and output `Tensor`s.
          Default value: None which means that the dtype implied by `times` is used.
        name: Python string. The name to give this op.
          Default value: `None` which maps to `milstein_sample`.
      """
      name = name or "milstein_sample"
      with tf.name_scope(name):
          if stratonovich_order <= 0:
              raise ValueError(
                  "`stratonovich_order` must be a positive integer."
              )
          times = tf.convert_to_tensor(times, dtype=dtype)
          if dtype is None:
              dtype = times.dtype
          if initial_state is None:
              initial_state = tf.zeros(dim, dtype=dtype)
          initial_state = tf.convert_to_tensor(
              initial_state, dtype=dtype, name="initial_state"
          )
          num_requested_times = tff_utils.get_shape(times)[0]
          # Create a time grid for the Milstein scheme.
          if num_time_steps is not None and time_step is not None:
              raise ValueError(
                  "Only one of either `num_time_steps` or `time_step` "
                  "should be defined but not both"
              )
          if time_step is None:
              if num_time_steps is None:
                  raise ValueError(
                      "Either `num_time_steps` or `time_step` should be "
                      "defined."
                  )
              num_time_steps = tf.convert_to_tensor(
                  num_time_steps, dtype=tf.int32, name="num_time_steps"
              )
              time_step = times[-1] / tf.cast(
                  num_time_steps, dtype=dtype
              )
          else:
              time_step = tf.convert_to_tensor(
                  time_step, dtype=dtype, name="time_step"
              )
          times, keep_mask, time_indices = utils.prepare_grid(
              times=times,
              time_step=time_step,
              num_time_steps=num_time_steps,
              dtype=dtype,
          )
          if watch_params is not None:
              watch_params = [
                  tf.convert_to_tensor(param, dtype=dtype)
                  for param in watch_params
              ]
          if grad_volatility_fn is None:
  
              def _grad_volatility_fn(
                  current_time, current_state, input_gradients
              ):
                  return gradient.fwd_gradient(
                      functools.partial(volatility_fn, current_time),
                      current_state,
                      input_gradients=input_gradients,
                      unconnected_gradients=tf.UnconnectedGradients.ZERO,
                  )
  
              grad_volatility_fn = _grad_volatility_fn
  
          input_gradients = None
          if dim > 1:
              input_gradients = tf.unstack(tf.eye(dim, dtype=dtype))
              input_gradients = [
                  tf.broadcast_to(start, [num_samples, dim])
                  for start in input_gradients
              ]
  
          return _sample(
              dim=dim,
              drift_fn=drift_fn,
              volatility_fn=volatility_fn,
              grad_volatility_fn=grad_volatility_fn,
              times=times,
              time_step=time_step,
              keep_mask=keep_mask,
              num_requested_times=num_requested_times,
              num_samples=num_samples,
              initial_state=initial_state,
              random_type=random_type,
              seed=seed,
              swap_memory=swap_memory,
              skip=skip,
              precompute_normal_draws=precompute_normal_draws,
              watch_params=watch_params,
              time_indices=time_indices,
              input_gradients=input_gradients,
              stratonovich_order=stratonovich_order,
              dtype=dtype,
          )
  
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 1 for calling a global function with example usage: sample() and returns: _sample, gradient