code: def _while_loop(
      *,
      steps_num,
      current_state,
      drift_fn,
      volatility_fn,
      wiener_mean,
      num_samples,
      times,
      dt,
      sqrt_dt,
      num_requested_times,
      keep_mask,
      swap_memory,
      random_type,
      seed,
      normal_draws,
      dtype,
  ):
      """Sample paths using tf.while_loop."""
      written_count = 0
      if (
          isinstance(num_requested_times, int)
          and num_requested_times == 1
      ):
          record_samples = False
          result = current_state
      else:
          # If more than one sample has to be recorded, create a TensorArray
          record_samples = True
          element_shape = current_state.shape
          result = tf.TensorArray(
              dtype=dtype,
              size=num_requested_times,
              element_shape=element_shape,
              clear_after_read=False,
          )
          # Include initial state, if necessary
          result = result.write(written_count, current_state)
      written_count += tf.cast(keep_mask[0], dtype=tf.int32)
  
      # Define sampling while_loop body function
      def cond_fn(i, written_count, *args):
          # It can happen that `times_grid[-1] > times[-1]` in which case we have
          # to terminate when `written_count` reaches `num_requested_times`
          del args
          return tf.math.logical_and(
              i < steps_num, written_count < num_requested_times
          )
  
      def step_fn(i, written_count, current_state, result):
          return _euler_step(
              i=i,
              written_count=written_count,
              current_state=current_state,
              result=result,
              drift_fn=drift_fn,
              volatility_fn=volatility_fn,
              wiener_mean=wiener_mean,
              num_samples=num_samples,
              times=times,
              dt=dt,
              sqrt_dt=sqrt_dt,
              keep_mask=keep_mask,
              random_type=random_type,
              seed=seed,
              normal_draws=normal_draws,
              record_samples=record_samples,
          )
  
      # Sample paths
      _, _, _, result = tf.while_loop(
          cond_fn,
          step_fn,
          (0, written_count, current_state, result),
          maximum_iterations=steps_num,
          swap_memory=swap_memory,
      )
      if not record_samples:
          # shape batch_shape + [num_samples, 1, dim]
          return tf.expand_dims(result, axis=-2)
      # Shape [num_time_points] + batch_shape + [num_samples, dim]
      result = result.stack()
      # transpose to shape batch_shape + [num_samples, num_time_points, dim]
      n = result.shape.rank
      perm = list(range(1, n - 1)) + [0, n - 1]
      return tf.transpose(result, perm)
  
file_path: tf-quant-finance/tf_quant_finance/models/euler_sampling.py
imports: from typing import Callable, List, Optional
  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 random
  from tf_quant_finance.models import utils
label_desc: Use image node_id 3 for calling a global function with example usage: _while_loop() and returns: tf, tf, _euler_step, tf