Python tensorflow.compat.v2.where() Examples

def _updated_cashflow(num_times, exercise_index, exercise_value,
                      expected_continuation, cashflow):
  """Revises the cashflow tensor where options will be exercised earlier."""
  do_exercise_bool = exercise_value > expected_continuation
  do_exercise = tf.cast(do_exercise_bool, exercise_value.dtype)
  # Shape [num_samples, payoff_dim]
  scaled_do_exercise = tf.where(do_exercise_bool, exercise_value,
                                tf.zeros_like(exercise_value))
  # This picks out the samples where we now wish to exercise.
  # Shape [num_samples, payoff_dim, 1]
  new_samp_masked = tf.expand_dims(scaled_do_exercise, axis=2)
  # This should be one on the current time step and zero otherwise.
  # This is an array with nonzero entries showing newly exercised payoffs.
  zeros = tf.zeros_like(cashflow)
  mask = tf.equal(tf.range(0, num_times), exercise_index - 1)
  new_cash = tf.where(mask, new_samp_masked, zeros)
  # Has shape [num_samples, payoff_dim, 1]
  old_mask = tf.expand_dims(1 - do_exercise, axis=2)
  mask = tf.range(0, num_times) >= exercise_index
  old_mask = tf.where(mask, old_mask, zeros)
  # Shape [num_samples, payoff_dim, num_times]
  old_cash = old_mask * cashflow
  return new_cash + old_cash 

def _should_stop(state, stopping_policy_fn):
  """Indicates whether the overall Brent search should continue.

  Args:
    state: A Python `_BrentSearchState` namedtuple.
    stopping_policy_fn: Python `callable` controlling the algorithm termination.

  Returns:
    A boolean value indicating whether the overall search should continue.
  """
  return tf.convert_to_tensor(
      stopping_policy_fn(state.finished), name="should_stop", dtype=tf.bool)


# This is a direct translation of the Brent root-finding method.
# Each operation is guarded by a call to `tf.where` to avoid performing
# unnecessary calculations. 

def business_days_between(self, from_dates, to_dates):
    """Calculates number of business between pairs of dates.

    For each pair, the initial date is included in the difference, and the final
    date is excluded. If the final date is the same or earlier than the initial
    date, zero is returned.

    Args:
      from_dates: `DateTensor` of initial dates.
      to_dates: `DateTensor` of final dates, should be broadcastable to
        `from_dates`.

    Returns:
       An int32 Tensor with the number of business days between the
       corresponding pairs of dates.
    """
    from_biz, from_is_bizday = self._to_biz_space(
        dt.convert_to_date_tensor(from_dates).ordinal())
    to_biz, to_is_bizday = self._to_biz_space(
        dt.convert_to_date_tensor(to_dates).ordinal())
    from_biz = tf.where(from_is_bizday, from_biz, from_biz + 1)
    to_biz = tf.where(to_is_bizday, to_biz, to_biz + 1)
    return tf.math.maximum(to_biz - from_biz, 0) 

def to_tensor(self):
    """Packs the dates into a single Tensor.

    The Tensor has shape `date_tensor.shape() + (3,)`, where the last dimension
    represents years, months and days, in this order.

    This can be convenient when the dates are the final result of a computation
    in the graph mode: a `tf.function` can return `date_tensor.to_tensor()`, or,
    if one uses `tf.compat.v1.Session`, they can call
    `session.run(date_tensor.to_tensor())`.

    Returns:
      A Tensor of shape `date_tensor.shape() + (3,)`.

    #### Example

    ```python
    dates = tff.datetime.dates_from_tuples([(2019, 1, 25), (2020, 3, 2)])
    dates.to_tensor()  # tf.Tensor with contents [[2019, 1, 25], [2020, 3, 2]].
    ```
    """
    return tf.stack((self.year(), self.month(), self.day()), axis=-1) 

def build_model(self, feature_ph_dict, labels_ph, params=None, context=None):
    """(DEPRECATED) Construct and return a Model.

    Args:
      feature_ph_dict: maps feature names to placeholders that will hold the
        corresponding inputs.
      labels_ph: a placeholder that will hold the target labels
      params: optional parameters to be passed to the config_* methods.
      context: if provided, use instead of building a fresh context

    Returns:
      a fully configured Model, where model.context is a
      freshly-built context produced by self.build_context().
    """
    model = Model()
    model.context = context or self.build_context(params=params)
    self.config_model(model, feature_ph_dict, labels_ph, params=params)
    self.check_model_completeness(model)
    return model 

def labels_of_top_ranked_predictions_in_batch(labels, predictions):
  """Applying tf.metrics.mean to this gives precision at 1.

  Args:
    labels: minibatch of dense 0/1 labels, shape [batch_size rows, num_classes]
    predictions: minibatch of predictions of the same shape

  Returns:
    one-dimension tensor top_labels, where top_labels[i]=1.0 iff the
    top-scoring prediction for batch element i has label 1.0
  """
  indices_of_top_preds = tf.cast(tf.argmax(input=predictions, axis=1), tf.int32)
  batch_size = tf.reduce_sum(input_tensor=tf.ones_like(indices_of_top_preds))
  row_indices = tf.range(batch_size)
  thresholded_labels = tf.where(labels > 0.0, tf.ones_like(labels),
                                tf.zeros_like(labels))
  label_indices_to_gather = tf.transpose(
      a=tf.stack([row_indices, indices_of_top_preds]))
  return tf.gather_nd(thresholded_labels, label_indices_to_gather) 

def one_hot_numpy_array(self, entity_name,
                          type_name):
    """A one-hot 1-by-N matrix encoding this entity.

    Args:
      entity_name: a string naming an entity.
      type_name: the string type name of the named entity

    Returns:
      A numpy array with shape (1,n) where n is the number of columns.

    Raises:
      EntityNameError: if entity_name does not map to a legal id.
    """
    index = self.get_id(entity_name, type_name)
    if index is None:
      raise EntityNameError(
          entity_name=entity_name,
          type_name=type_name,
          message='Cannot make a one-hot vector')
    result = self.zeros_numpy_array(type_name, True)
    result[0, index] = 1.0
    return result

  # schema operations 

def nonneg_softmax(expr,
                   replace_nonpositives = -10):
  """A softmax operator that is appropriate for NQL outputs.

  NeuralQueryExpressions often evaluate to sparse vectors of small, nonnegative
  values. Softmax for those is dominated by zeros, so this is a fix.  This also
  fixes the problem that minibatches for NQL are one example per column, not one
  example per row.

  Args:
    expr: a Tensorflow expression for some predicted values.
    replace_nonpositives: will replace zeros with this value before computing
      softmax.

  Returns:
    Tensorflow expression for softmax.
  """
  if replace_nonpositives != 0.0:
    ones = tf.ones(tf.shape(input=expr), tf.float32)
    expr = tf.where(expr > 0.0, expr, ones * replace_nonpositives)
  return tf.nn.softmax(expr) 

def nonneg_crossentropy(expr, target):
  """A cross entropy operator that is appropriate for NQL outputs.

  Query expressions often evaluate to sparse vectors.  This evaluates cross
  entropy safely.

  Args:
    expr: a Tensorflow expression for some predicted values.
    target: a Tensorflow expression for target values.

  Returns:
    Tensorflow expression for cross entropy.
  """
  expr_replacing_0_with_1 = \
     tf.where(expr > 0, expr, tf.ones(tf.shape(input=expr), tf.float32))
  cross_entropies = tf.reduce_sum(
      input_tensor=-target * tf.math.log(expr_replacing_0_with_1), axis=1)
  return tf.reduce_mean(input_tensor=cross_entropies, axis=0) 

def __call__(self, multi_objectives: tf.Tensor) -> tf.Tensor:
    """Returns a single reward by scalarizing multiple objectives.

    Args:
      multi_objectives: A `Tensor` of shape [batch_size, number_of_objectives],
        where each column represents an objective.

    Returns: A `Tensor` of shape [batch_size] representing scalarized rewards.

    Raises:
      ValueError: if `multi_objectives.shape.rank != 2`.
      ValueError: if
        `multi_objectives.shape.dims[1] != self._num_of_objectives`.
    """
    if multi_objectives.shape.rank != 2:
      raise ValueError('The rank of the input should be 2, but is {}'.format(
          multi_objectives.shape.rank))
    if multi_objectives.shape.dims[1] != self._num_of_objectives:
      raise ValueError(
          'The number of input objectives should be {}, but is {}.'.format(
              self._num_of_objectives, multi_objectives.shape.dims[1]))
    return self.call(multi_objectives)

  # Subclasses must implement these methods. 

def _get_forward_rate(self, valuation_date, market):
    """Returns the relevant forward rates from the market data."""

    forward_rates = market.reference_curve.get_forward_rate(
        self._accrual_start_dates,
        self._accrual_end_dates,
        self._daycount_fractions)

    forward_rates = tf.where(self._daycount_fractions > 0.0, forward_rates,
                             tf.zeros_like(forward_rates))

    libor_rate = rc.get_rate_index(
        market, self._start_date, rc.RateIndexType.LIBOR, dtype=self._dtype)
    libor_rate = tf.repeat(
        tf.convert_to_tensor(libor_rate, dtype=self._dtype), self._num_caplets)
    forward_rates = tf.where(
        self._accrual_end_dates < valuation_date,
        tf.constant(0., dtype=self._dtype),
        tf.where(self._accrual_start_dates < valuation_date, libor_rate,
                 forward_rates))

    return forward_rates 

def _log_prob_with_logsf_and_logcdf(self, y):
    """Compute log_prob(y) using log survival_function and cdf together."""
    # There are two options that would be equal if we had infinite precision:
    # Log[ sf(y - .5) - sf(y + .5) ]
    #   = Log[ exp{logsf(y - .5)} - exp{logsf(y + .5)} ]
    # Log[ cdf(y + .5) - cdf(y - .5) ]
    #   = Log[ exp{logcdf(y + .5)} - exp{logcdf(y - .5)} ]
    logsf_y_plus = self.base.log_survival_function(y + .5)
    logsf_y_minus = self.base.log_survival_function(y - .5)
    logcdf_y_plus = self.base.log_cdf(y + .5)
    logcdf_y_minus = self.base.log_cdf(y - .5)

    # Important:  Here we use select in a way such that no input is inf, this
    # prevents the troublesome case where the output of select can be finite,
    # but the output of grad(select) will be NaN.

    # In either case, we are doing Log[ exp{big} - exp{small} ]
    # We want to use the sf items precisely when we are on the right side of the
    # median, which occurs when logsf_y < logcdf_y.
    condition = logsf_y_plus < logcdf_y_plus
    big = tf.where(condition, logsf_y_minus, logcdf_y_plus)
    small = tf.where(condition, logsf_y_plus, logcdf_y_minus)
    return _logsum_expbig_minus_expsmall(big, small) 

def select_actor_action(self, env_output, agent_output):
    oracle_next_action = env_output.observation[constants.ORACLE_NEXT_ACTION]
    oracle_next_action_indices = tf.where(
        tf.equal(env_output.observation[constants.CONN_IDS],
                 oracle_next_action))
    oracle_next_action_idx = tf.reduce_min(oracle_next_action_indices)
    assert self._mode, 'mode must be set.'
    if self._mode == 'train':
      if self._loss_type == common.CE_LOSS:
        # This is teacher-forcing mode, so choose action same as oracle action.
        action_idx = oracle_next_action_idx
      elif self._loss_type == common.AC_LOSS:
        # Choose next pano from probability distribution over next panos
        action_idx = tfp.distributions.Categorical(
            logits=agent_output.policy_logits).sample()
      else:
        raise ValueError('Unsupported loss type {}'.format(self._loss_type))
    else:
      # In non-train modes, choose greedily.
      action_idx = tf.argmax(agent_output.policy_logits, axis=-1)
    action_val = env_output.observation[constants.CONN_IDS][action_idx]
    return common.ActorAction(
        chosen_action_idx=int(action_idx.numpy()),
        oracle_next_action_idx=int(oracle_next_action_idx.numpy())), int(
            action_val.numpy()) 

def select_actor_action(self, env_output, agent_output):
    oracle_next_action = env_output.observation[constants.ORACLE_NEXT_ACTION]
    oracle_next_action_indices = tf.where(
        tf.equal(env_output.observation[constants.CONN_IDS],
                 oracle_next_action))
    oracle_next_action_idx = tf.reduce_min(oracle_next_action_indices)
    if self._loss_type == common.CE_LOSS:
      # This is teacher-forcing mode, so choose action same as oracle action.
      action_idx = oracle_next_action_idx
    elif self._loss_type == common.AC_LOSS:
      # Choose next pano from probability distribution over next panos
      action_idx = tfp.distributions.Categorical(
          logits=agent_output.policy_logits).sample()
    else:
      raise ValueError('Unsupported loss type {}'.format(self._loss_type))
    action_val = env_output.observation[constants.CONN_IDS][action_idx]
    return common.ActorAction(
        chosen_action_idx=int(action_idx.numpy()),
        oracle_next_action_idx=int(oracle_next_action_idx.numpy())), int(
            action_val.numpy()) 

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 select_actor_action(self, env_output, agent_output):
    oracle_next_action = env_output.observation[constants.ORACLE_NEXT_ACTION]
    oracle_next_action_indices = tf.where(
        tf.equal(env_output.observation[constants.CONN_IDS],
                 oracle_next_action))
    oracle_next_action_idx = tf.reduce_min(oracle_next_action_indices)
    assert self._mode, 'mode must be set.'
    if self._mode == 'train':
      if self._loss_type == common.CE_LOSS:
        # This is teacher-forcing mode, so choose action same as oracle action.
        action_idx = oracle_next_action_idx
      elif self._loss_type == common.AC_LOSS:
        # Choose next pano from probability distribution over next panos
        action_idx = tfp.distributions.Categorical(
            logits=agent_output.policy_logits).sample()
      else:
        raise ValueError('Unsupported loss type {}'.format(self._loss_type))
    else:
      # In non-train modes, choose greedily.
      action_idx = tf.argmax(agent_output.policy_logits, axis=-1)
    action_val = env_output.observation[constants.CONN_IDS][action_idx]
    return common.ActorAction(
        chosen_action_idx=int(action_idx.numpy()),
        oracle_next_action_idx=int(oracle_next_action_idx.numpy())), int(
            action_val.numpy()) 

def triu(m, k=0):  # pylint: disable=missing-docstring
  m = asarray(m).data
  m_shape = m.shape.as_list()

  if len(m_shape) < 2:
    raise ValueError('Argument to triu must have rank at least 2')

  if m_shape[-1] is None or m_shape[-2] is None:
    raise ValueError('Currently, the last two dimensions of the input array '
                     'need to be known.')

  z = tf.constant(0, m.dtype)

  mask = tri(*m_shape[-2:], k=k - 1, dtype=bool)
  return utils.tensor_to_ndarray(
      tf.where(tf.broadcast_to(mask, tf.shape(m)), z, m)) 

def tril(m, k=0):  # pylint: disable=missing-docstring
  m = asarray(m).data
  m_shape = m.shape.as_list()

  if len(m_shape) < 2:
    raise ValueError('Argument to tril must have rank at least 2')

  if m_shape[-1] is None or m_shape[-2] is None:
    raise ValueError('Currently, the last two dimensions of the input array '
                     'need to be known.')

  z = tf.constant(0, m.dtype)

  mask = tri(*m_shape[-2:], k=k, dtype=bool)
  return utils.tensor_to_ndarray(
      tf.where(tf.broadcast_to(mask, tf.shape(m)), m, z)) 

def stack_nested_tensors(list_of_nests):
  """Stack a list of nested tensors.

  Args:
    list_of_nests: A list of nested tensors (or numpy arrays) of the same
      shape/structure.

  Returns:
    A nested array containing batched items, where each batched item is obtained
    by stacking corresponding items from the list of nested_arrays.
  """


  def stack_tensor(*tensors):
    result = [tf.convert_to_tensor(t) for t in tensors]
    return tf.stack(result)

  return tf.nest.map_structure(stack_tensor, *list_of_nests) 

def as_nql(
      self,
      expr,
      type_name,
      provenance = None
  ):
    """Convert a Tensorflow expression to a NeuralQueryExpression.

    Args:
      expr: a Tensorflow expression which evaluates to a dense matrix of size
        N*M, where M is the minibatch size, and N is the max id for the named
        type.  Alternatively, expr can be a NeuralQueryExpression, which will be
        typechecked and returned unchanged
      type_name: string naming a type
      provenance: if provided, an NQExprProvenance object

    Returns:
      A NeuralQueryExpression

    Raises:
      TypeNameError: If type_name is not a known type.
      TypeCompatibilityError: If expression is not of type_name type.
      TypeError: If the factory method returned an unexpected type.
    """
    if isinstance(expr, NeuralQueryExpression):
      if expr.type_name != type_name:
        raise TypeCompatibilityError(expr.type_name, type_name, 'as_nql')
      return expr
    else:
      tf_expr = expr
      result = self.expression_factory_class(
          self, tf_expr, type_name, provenance=provenance)
      if not isinstance(result, NeuralQueryExpression):
        raise TypeError(
            'Factory returned %r instead of a NeuralQueryExpression:' %
            type(result))
      return result 

def _prob_with_sf_and_cdf(self, y):
    # There are two options that would be equal if we had infinite precision:
    # sf(y - .5) - sf(y + .5)
    # cdf(y + .5) - cdf(y - .5)
    sf_y_plus = self.base.survival_function(y + .5)
    sf_y_minus = self.base.survival_function(y - .5)
    cdf_y_plus = self.base.cdf(y + .5)
    cdf_y_minus = self.base.cdf(y - .5)

    # sf_prob has greater precision iff we're on the right side of the median.
    return tf.where(
        sf_y_plus < cdf_y_plus,
        sf_y_minus - sf_y_plus, cdf_y_plus - cdf_y_minus) 

def call(self, multi_objectives: tf.Tensor) -> tf.Tensor:
    transformed_objectives = tf.maximum(
        multi_objectives * self._slopes + self._offsets, 0)
    nonzero_mask = tf.broadcast_to(
        tf.cast(tf.abs(self._direction) >= self.ALMOST_ZERO, dtype=tf.bool),
        multi_objectives.shape)
    return tf.reduce_min(
        tf.where(nonzero_mask, transformed_objectives / self._direction,
                 multi_objectives.dtype.max),
        axis=1) 

def select_actor_action(self, env_output, agent_output):
    # Agent_output is unused here.
    oracle_next_action = env_output.observation[constants.ORACLE_NEXT_ACTION]
    oracle_next_action_indices = tf.where(
        tf.equal(env_output.observation[constants.CONN_IDS],
                 oracle_next_action))
    oracle_next_action_idx = tf.reduce_min(oracle_next_action_indices)
    assert self._mode, 'mode must be set.'
    action_idx = oracle_next_action_idx
    action_val = env_output.observation[constants.CONN_IDS][action_idx]
    return common.ActorAction(
        chosen_action_idx=int(action_idx.numpy()),
        oracle_next_action_idx=int(oracle_next_action_idx.numpy())), int(
            action_val) 

def weighted_by(self, weight_rel_name,
                  category_name):
    """Weight elements in set by confidence in some relationship holding.

    Specifically, weight all elements x in a set by the confidence
    given to the trip weight_rel_name(x,category_name).
    is equivalent to

    x * x.jump_to(category_name,category_type).follow(weight_rel_name,-1)

    where category_type is the appropriate type for, is determined by the
    declaration for weight_rel_name.

    Args:
      weight_rel_name: a string naming a declared KG relation
      category_name: an entity name in the KG

    Returns:
      A NeuralQueryExpression
    """
    self._check_type_compatibility(
        self.context.get_domain(weight_rel_name), self._type_name,
        'weighted_by')
    category_type = self.context.get_range(weight_rel_name)
    with tf.name_scope('weighted_by_%s_%s' % (weight_rel_name, category_name)):
      weight_vector = self.context.one(category_name, category_type).follow(
          weight_rel_name, -1)
      return self.__mul__(weight_vector) 

def _apply_roll_biz_space(self, date_tensor, biz_days, is_bizday,
                            roll_convention):
    """Applies roll in business day space."""
    if roll_convention == constants.BusinessDayConvention.NONE:
      # If no business convention is specified, return the current business
      # day.
      return biz_days

    if roll_convention == constants.BusinessDayConvention.FOLLOWING:
      return tf.where(is_bizday, biz_days, biz_days + 1)

    if roll_convention == constants.BusinessDayConvention.PRECEDING:
      return biz_days

    if roll_convention == constants.BusinessDayConvention.MODIFIED_FOLLOWING:
      maybe_prev_biz_day = biz_days
      maybe_next_biz_day = tf.where(is_bizday, biz_days, biz_days + 1)
      maybe_next_biz_ordinal = self._from_biz_space(maybe_next_biz_day)
      take_previous = tf.not_equal(
          _get_month(maybe_next_biz_ordinal), date_tensor.month())
      return tf.where(take_previous, maybe_prev_biz_day, maybe_next_biz_day)

    if roll_convention == constants.BusinessDayConvention.MODIFIED_PRECEDING:
      maybe_prev_biz_day = biz_days
      maybe_next_biz_day = tf.where(is_bizday, biz_days, biz_days + 1)
      maybe_prev_biz_ordinal = self._from_biz_space(maybe_prev_biz_day)
      take_next = tf.not_equal(
          _get_month(maybe_prev_biz_ordinal), date_tensor.month())
      return tf.where(take_next, maybe_next_biz_day, maybe_prev_biz_day)

    raise ValueError('Unsupported roll convention: {}'.format(roll_convention)) 

def _compute_bus_day_ordinals_table(self):
    """Computes and caches rolled business day ordinals table."""
    if self._table_cache.bus_day_ordinals is not None:
      return self._table_cache.bus_day_ordinals

    is_bus_day_table = self._compute_is_bus_day_table()
    with tf.init_scope():
      bus_day_ordinals_table = (
          tf.cast(tf.where(is_bus_day_table)[:, 0], tf.int32) +
          self._ordinal_offset - 1)
      self._table_cache.bus_day_ordinals = bus_day_ordinals_table
    return bus_day_ordinals_table 

def _region_3(g1plus2g0, g0plus2g1, g0, g1, x):
  """Computes conditional and value for points in region 3."""
  del g1plus2g0
  # Reference: Eq. 30, 31 in Ref [2]
  is_region_3 = (((g1 <= 0) & (g0plus2g1 > 0)) | ((g1 >= 0) & (g0plus2g1 < 0)))
  eta = 3 * g1 / (g1 - g0)
  x_cap = tf.math.minimum(x, eta)
  ratio = (eta - x_cap) / eta
  # Replace NaN values (corresponding to g1 == 0) with zeros.
  ratio = tf.where(tf.math.is_nan(ratio), tf.zeros_like(ratio), ratio)
  region_3_value = g1 + (g0 - g1) * tf.math.square(ratio)
  integrated_value = g1 * x + eta * (g0 - g1) / 3 * (1 - ratio**3)
  return is_region_3, region_3_value, integrated_value 

def _interpolate_adjacent(times, values, name=None):
  """Interpolates linearly between adjacent values.

  Suppose `times` are `[t_1, t_2, ..., t_n]` an array of length `n` and
  values are `[f_1, ... f_n]` of length `n`. This function associates
  each of the values to the midpoint of the interval i.e. `f_i` is associated
  to the midpoint of the interval `[t_i, t_{i+1}]`. Then it calculates the
  values at the interval boundaries by linearly interpolating between adjacent
  intervals. The first interval is considered to be `[0, t_1]`. The values at
  the endpoints (i.e. result[0] and result[n]) are computed as follows:
  `result[0] = values[0] - 0.5 * (result[1] - values[0])` and
  `result[n] = values[n-1] - 0.5 * (result[n-1] - values[n-1])`.
  The rationale for these specific values is discussed in Ref. [1].

  Args:
    times: A rank 1 `Tensor` of real dtype. The times at which the interpolated
      values are to be computed. The values in the array should be positive and
      monotonically increasing.
    values: A rank 1 `Tensor` of the same dtype and shape as `times`. The values
      assigned to the midpoints of the time intervals.
    name: Python `str` name prefixed to Ops created by this class.
      Default value: None which is mapped to the default name
        'interpolate_adjacent'.

  Returns:
    interval_values: The values interpolated from the supplied midpoint values
      as described above. A `Tensor` of the same dtype as `values` but shape
      `[n+1]` where `[n]` is the shape of `values`. The `i`th component of the
      is the value associated to the time point `t_{i+1}` with `t_0 = 0`.
  """
  with tf.compat.v1.name_scope(
      name, default_name='interpolate_adjacent', values=[times, values]):
    dt1 = diff(times, order=1, exclusive=False)
    dt2 = diff(times, order=2, exclusive=False)[1:]
    weight_right = dt1[:-1] / dt2
    weight_left = dt1[1:] / dt2
    interior_values = weight_right * values[1:] + weight_left * values[:-1]
    value_0 = values[0] - 0.5 * (interior_values[0] - values[0])
    value_n = values[-1] - 0.5 * (interior_values[-1] - values[-1])
    return tf.concat([[value_0], interior_values, [value_n]], axis=0) 

def _lsm_loop_body(sample_paths, exercise_times, discount_factors, payoff_fn,
                   basis_fn, num_times, exercise_index, cashflow):
  """Finds the optimal exercise point and updates `cashflow`."""

  # Index of the sample path that the exercise index maps to.
  time_index = exercise_times[exercise_index - 1]
  # Calculate the payoff of each path if exercised now.
  # Shape [num_samples, payoff_dim]
  exercise_value = payoff_fn(sample_paths, time_index)
  # Present value of hanging on to the options (using future information).
  # Shape `[num_samples, payoff_dim]`
  continuation_value = _continuation_value_fn(cashflow, discount_factors,
                                              exercise_index)
  # Create a design matrix for regression based on the sample paths.
  # Shape `[num_samples, basis_size]`
  design = basis_fn(sample_paths, time_index)

  # Expected present value of hanging on the options.
  expected_continuation = _expected_exercise_fn(design, continuation_value,
                                                exercise_value)
  # Update the cashflow matrix to reflect where earlier exercise is optimal.
  # Shape `[num_samples, payoff_dim, num_exercise]`.
  rev_cash = _updated_cashflow(num_times, exercise_index, exercise_value,
                               expected_continuation,
                               cashflow)
  return LsmLoopVars(exercise_index=exercise_index - 1, cashflow=rev_cash) 

def _analytic_valuation(discount_rate_fn, model, strikes, expiries, maturities,
                        dim, is_call_options):
  """Performs analytic valuation."""

  discount_factor_expiry = tf.math.exp(
      -discount_rate_fn(expiries) * expiries)
  input_shape = expiries.shape
  variance = _bond_option_variance(
      model, tf.reshape(expiries, shape=[-1]), tf.reshape(maturities, [-1]),
      dim)
  variance = tf.reshape(variance, [dim] + input_shape)
  discount_factor_maturity = tf.math.exp(-discount_rate_fn(maturities) *
                                         maturities)
  forward_bond_price = discount_factor_maturity / discount_factor_expiry
  sqrt_variance = tf.math.sqrt(variance)
  d1 = (tf.expand_dims(tf.math.log(forward_bond_price / strikes), axis=0) +
        0.5 * variance) / sqrt_variance
  d2 = d1 - tf.math.sqrt(variance)
  option_value_call = (
      tf.expand_dims(discount_factor_maturity, axis=0) * _ncdf(d1) -
      tf.expand_dims(strikes * discount_factor_expiry, axis=0) * _ncdf(d2))
  option_value_put = (
      tf.expand_dims(strikes * discount_factor_expiry, axis=0) * _ncdf(-d2)
      - tf.expand_dims(discount_factor_maturity, axis=0) * _ncdf(-d1))

  is_call_options = tf.broadcast_to(is_call_options, [dim] + strikes.shape)
  option_value = tf.where(is_call_options, option_value_call,
                          option_value_put)

  # Make `dim` as the last dimension and return.
  return tf.transpose(
      option_value,
      perm=[i for i in range(1, len(option_value.shape.as_list()))] + [0])


# TODO(b/158501671): Clean-up this implementation.