Python tensorflow.compat.v2.while_loop() Examples

def _tf_gcd(x1, x2):
  def _gcd_cond_fn(x1, x2):
    return tf.reduce_any(x2 != 0)
  def _gcd_body_fn(x1, x2):
    # tf.math.mod will raise an error when any element of x2 is 0. To avoid
    # that, we change those zeros to ones. Their values don't matter because
    # they won't be used.
    x2_safe = tf.where(x2 != 0, x2, tf.constant(1, x2.dtype))
    x1, x2 = (tf.where(x2 != 0, x2, x1),
              tf.where(x2 != 0, tf.math.mod(x1, x2_safe),
                       tf.constant(0, x2.dtype)))
    return (tf.where(x1 < x2, x2, x1), tf.where(x1 < x2, x1, x2))
  if (not np.issubdtype(x1.dtype.as_numpy_dtype, np.integer) or
      not np.issubdtype(x2.dtype.as_numpy_dtype, np.integer)):
    raise ValueError("Arguments to gcd must be integers.")
  shape = tf.broadcast_static_shape(x1.shape, x2.shape)
  x1 = tf.broadcast_to(x1, shape)
  x2 = tf.broadcast_to(x2, shape)
  gcd, _ = tf.while_loop(_gcd_cond_fn, _gcd_body_fn,
                         (tf.math.abs(x1), tf.math.abs(x2)))
  return gcd 

def _while_loop(*, dim, steps_num, current_state,
                drift_fn, volatility_fn, wiener_mean,
                num_samples, times, dt, sqrt_dt, time_step, num_requested_times,
                keep_mask, swap_memory, random_type, seed, normal_draws):
  """Smaple paths using tf.while_loop."""
  cond_fn = lambda i, *args: i < steps_num
  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)
  maximum_iterations = (tf.cast(1. / time_step, dtype=tf.int32)
                        + tf.size(times))
  result = tf.zeros((num_samples, num_requested_times, dim),
                    dtype=current_state.dtype)
  _, _, _, result = tf.while_loop(
      cond_fn, step_fn, (0, 0, current_state, result),
      maximum_iterations=maximum_iterations,
      swap_memory=swap_memory)
  return result 

def create_tf_while_loop_fn(step_fn):
  """Create a multiple steps function driven by tf.while_loop on the host.

  Args:
    step_fn: A function which takes `iterator` as input.

  Returns:
    A callable defined as the `loop_fn` defination below.
  """

  @tf.function
  def loop_fn(iterator, num_steps):
    """A loop function with multiple steps.

    Args:
      iterator: A nested structure of tf.data `Iterator` or
        `DistributedIterator`.
      num_steps: The number of steps in the loop. Must be a tf.Tensor.
    """
    if not isinstance(num_steps, tf.Tensor):
      raise ValueError("`num_steps` should be an `tf.Tensor`. Python object "
                       "may cause retracing.")

    for _ in tf.range(num_steps):
      step_fn(iterator)

  return loop_fn 

def sample_paths(self,
                   times,
                   num_samples=1,
                   initial_state=None,
                   random_type=None,
                   seed=None,
                   swap_memory=True,
                   name=None,
                   **kwargs):
    """Returns a sample of paths from the process.

    Generates samples of paths from the process at the specified time points.

    Args:
      times: Rank 1 `Tensor` of increasing positive real values. The times at
        which the path points are to be evaluated.
      num_samples: Positive scalar `int`. The number of paths to draw.
      initial_state: `Tensor` of shape `[dim]`. The initial state of the
        process.
        Default value: None which maps to a zero initial state.
      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: Python `int`. The random seed to use. If not supplied, no seed is
        set.
      swap_memory: Whether GPU-CPU memory swap is enabled for this op. See
        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.
      name: str. The name to give this op. If not supplied, default name of
        `sample_paths` is used.
      **kwargs: parameters, specific to Euler schema: `grid_step` is rank 0 real
        `Tensor` - maximal distance between points in grid in Euler schema. Note
        that Euler sampling is only used if it is not possible to do exact
        sampling because total drift or total covariance are unavailable.

    Returns:
     A real `Tensor` of shape [num_samples, k, n] where `k` is the size of the
        `times`, `n` is the dimension of the process.
    """
    if self._total_drift_fn is None or self._total_covariance_fn is None:
      return super(BrownianMotion, self).sample_paths(
          times,
          num_samples=num_samples,
          initial_state=initial_state,
          random_type=random_type,
          seed=seed,
          name=name,
          **kwargs)

    default_name = self._name + '_sample_path'
    with tf.compat.v1.name_scope(
        name, default_name=default_name, values=[times, initial_state]):
      end_times = tf.convert_to_tensor(times, dtype=self.dtype())
      start_times = tf.concat(
          [tf.zeros([1], dtype=end_times.dtype), end_times[:-1]], axis=0)
      paths = self._exact_sampling(end_times, start_times, num_samples,
                                   initial_state, random_type, seed)
      if initial_state is not None:
        return paths + initial_state
      return paths 

def sample_paths(self,
                   times,
                   num_samples=1,
                   initial_state=None,
                   random_type=None,
                   seed=None,
                   swap_memory=True,
                   name=None,
                   **kwargs):
    """Returns a sample of paths from the process.

    The default implementation uses Euler schema. However, for particular types
    of Ito processes more efficient schemes can be used.

    Args:
      times: Rank 1 `Tensor` of increasing positive real values. The times at
        which the path points are to be evaluated.
      num_samples: Positive scalar `int`. The number of paths to draw.
      initial_state: `Tensor` of shape `[dim]`. The initial state of the
        process.
        Default value: None which maps to a zero initial state.
      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: Python `int`. The random seed to use. If not supplied, no seed is
        set.
      swap_memory: Whether GPU-CPU memory swap is enabled for this op. See
        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.
      name: str. The name to give this op. If not supplied, default name of
        `sample_paths` is used.
      **kwargs: parameters, specific to Euler schema: `grid_step` is rank 0 real
        `Tensor` - maximal distance between points in grid in Euler schema.

    Returns:
     A real `Tensor` of shape [num_samples, k, n] where `k` is the size of the
        `times`, `n` is the dimension of the process.
    """
    if self.drift_fn() is None or self.volatility_fn() is None:
      raise NotImplementedError(
          'In order to use Euler scheme, both drift_fn and volatility_fn '
          'should be provided.')
    default_name = self.name() + '_sample_paths'
    with tf.compat.v1.name_scope(
        name, default_name=default_name, values=[times, initial_state]):
      if initial_state is None:
        initial_state = tf.zeros(self._dim, dtype=self._dtype)
      times = tf.convert_to_tensor(times, dtype=self._dtype)
      initial_state = tf.convert_to_tensor(
          initial_state, dtype=self._dtype, name='initial_state')
      num_requested_times = tf.shape(times)[-1]
      grid_step = kwargs['grid_step']
      times, keep_mask = self._prepare_grid(times, grid_step)
      return self._sample_paths(times, grid_step, keep_mask,
                                num_requested_times, num_samples, initial_state,
                                random_type, seed, swap_memory) 

def _sample_paths(self, times, grid_step, keep_mask, num_requested_times,
                    num_samples, initial_state, random_type, seed, swap_memory):
    """Returns a sample of paths from the process."""
    dt = times[1:] - times[:-1]
    sqrt_dt = tf.sqrt(dt)
    current_state = initial_state + tf.zeros(
        [num_samples, self.dim()], dtype=initial_state.dtype)
    steps_num = tf.shape(dt)[-1]
    wiener_mean = tf.zeros((self.dim(), 1), dtype=self._dtype)

    cond_fn = lambda i, *args: i < steps_num

    def step_fn(i, written_count, current_state, result):
      """Performs one step of Euler scheme."""
      current_time = times[i + 1]
      dw = random_ops.mv_normal_sample((num_samples,),
                                       mean=wiener_mean,
                                       random_type=random_type,
                                       seed=seed)
      dw = dw * sqrt_dt[i]
      dt_inc = dt[i] * self.drift_fn()(current_time, current_state)  # pylint: disable=not-callable
      dw_inc = tf.squeeze(
          tf.matmul(self.volatility_fn()(current_time, current_state), dw), -1)  # pylint: disable=not-callable
      next_state = current_state + dt_inc + dw_inc

      def write_next_state_to_result():
        # Replace result[:, written_count, :] with next_state.
        one_hot = tf.one_hot(written_count, depth=num_requested_times)
        mask = tf.expand_dims(one_hot > 0, axis=-1)
        return tf.where(mask, tf.expand_dims(next_state, axis=1), result)

      # Keep only states for times requested by user.
      result = tf.cond(keep_mask[i + 1],
                       write_next_state_to_result,
                       lambda: result)
      written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
      return i + 1, written_count, next_state, result

    # Maximum number iterations is passed to the while loop below. It improves
    # performance of the while loop on a GPU and is needed for XLA-compilation
    # comptatiblity
    maximum_iterations = (
        tf.cast(1. / grid_step, dtype=tf.int32) + tf.size(times))
    result = tf.zeros((num_samples, num_requested_times, self.dim()))
    _, _, _, result = tf.compat.v1.while_loop(
        cond_fn,
        step_fn, (0, 0, current_state, result),
        maximum_iterations=maximum_iterations,
        swap_memory=swap_memory)

    return result 

def _sample(*, dim, drift_fn, 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,
            dtype):
  """Returns a sample of paths from the process using Euler 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 whith 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):
    normal_draws = utils.generate_mc_normal_draws(
        num_normal_draws=dim, num_time_steps=steps_num,
        num_sample_paths=num_samples, random_type=random_type,
        dtype=dtype, seed=seed, skip=skip)
    wiener_mean = None
  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
  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, 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)
  else:
    # Use custom for_loop if `watch_params` is specified
    return _for_loop(
        steps_num=steps_num, current_state=current_state,
        drift_fn=drift_fn, volatility_fn=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) 

def estimate_tails(func, target, shape, dtype):
  """Estimates approximate tail quantiles.

  This runs a simple Adam iteration to determine tail quantiles. The
  objective is to find an `x` such that:
  ```
  func(x) == target
  ```
  For instance, if `func` is a CDF and the target is a quantile value, this
  would find the approximate location of that quantile. Note that `func` is
  assumed to be monotonic. When each tail estimate has passed the optimal value
  of `x`, the algorithm does 10 additional iterations and then stops.

  This operation is vectorized. The tensor shape of `x` is given by `shape`, and
  `target` must have a shape that is broadcastable to the output of `func(x)`.

  Arguments:
    func: A callable that computes cumulative distribution function, survival
      function, or similar.
    target: The desired target value.
    shape: The shape of the `tf.Tensor` representing `x`.
    dtype: The `tf.dtypes.Dtype` of the computation (and the return value).

  Returns:
    A `tf.Tensor` representing the solution (`x`).
  """
  with tf.name_scope("estimate_tails"):
    dtype = tf.as_dtype(dtype)
    shape = tf.convert_to_tensor(shape, tf.int32)
    target = tf.convert_to_tensor(target, dtype)

    def loop_cond(tails, m, v, count):
      del tails, m, v  # unused
      return tf.reduce_min(count) < 10

    def loop_body(tails, m, v, count):
      with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(tails)
        loss = abs(func(tails) - target)
      grad = tape.gradient(loss, tails)
      m = .5 * m + .5 * grad  # Adam mean estimate.
      v = .9 * v + .1 * tf.square(grad)  # Adam variance estimate.
      tails -= .5 * m / (tf.sqrt(v) + 1e-7)
      # Start counting when the gradient flips sign (note that this assumes
      # `tails` is initialized to zero).
      count = tf.where(
          tf.math.logical_or(count > 0, tails * grad > 0),
          count + 1, count)
      return tails, m, v, count

    init_tails = tf.zeros(shape, dtype=dtype)
    init_m = tf.zeros(shape, dtype=dtype)
    init_v = tf.ones(shape, dtype=dtype)
    init_count = tf.zeros(shape, dtype=tf.int32)
    return tf.while_loop(
        loop_cond, loop_body, (init_tails, init_m, init_v, init_count),
        back_prop=False)[0]