parameters: (have_log_loss=False)
return_statement: return vimco_advantage_fn
docstring: VIMCO (Variational Inference for Monte Carlo Objectives) baseline.
  
    Implements VIMCO baseline from the article of the same name:
  
    https://arxiv.org/pdf/1602.06725v2.pdf
  
    Given a `loss` tensor (containing non-negative probabilities or ratios),
    calculates the advantage VIMCO advantage via Eq. 9 of the above paper.
  
    The tensor `loss` should be shaped `[n, ...]`, with rank at least 1.  Here,
    the first axis is considered the single sampling dimension and `n` must
    be at least 2.  Specifically, the `StochasticTensor` is assumed to have
    used the `SampleValue(n)` value type with `n > 1`.
  
    Args:
      have_log_loss: Python `Boolean`.  If `True`, the loss is assumed to be the
        log loss.  If `False` (the default), it is assumed to be a nonnegative
        probability or probability ratio.
  
    Returns:
      Callable baseline function that takes the `StochasticTensor` (unused) and
      the downstream `loss`, and returns the VIMCO baseline for the loss.
docstring_summary: VIMCO (Variational Inference for Monte Carlo Objectives) baseline.
docstring_tokens: VIMCO ( Variational Inference for Monte Carlo Objectives ) baseline .
function: def get_vimco_advantage_fn(have_log_loss=False):
    """VIMCO (Variational Inference for Monte Carlo Objectives) baseline.
  
    Implements VIMCO baseline from the article of the same name:
  
    https://arxiv.org/pdf/1602.06725v2.pdf
  
    Given a `loss` tensor (containing non-negative probabilities or ratios),
    calculates the advantage VIMCO advantage via Eq. 9 of the above paper.
  
    The tensor `loss` should be shaped `[n, ...]`, with rank at least 1.  Here,
    the first axis is considered the single sampling dimension and `n` must
    be at least 2.  Specifically, the `StochasticTensor` is assumed to have
    used the `SampleValue(n)` value type with `n > 1`.
  
    Args:
      have_log_loss: Python `Boolean`.  If `True`, the loss is assumed to be the
        log loss.  If `False` (the default), it is assumed to be a nonnegative
        probability or probability ratio.
  
    Returns:
      Callable baseline function that takes the `StochasticTensor` (unused) and
      the downstream `loss`, and returns the VIMCO baseline for the loss.
    """
    def vimco_advantage_fn(_, loss, name=None):
      """Internal VIMCO function.
  
      Args:
        _: ignored `StochasticTensor`.
        loss: The loss `Tensor`.
        name: Python string, the name scope to use.
  
      Returns:
        The advantage `Tensor`.
      """
      with ops.name_scope(name, "VIMCOAdvantage", values=[loss]):
        loss = ops.convert_to_tensor(loss)
        loss_shape = loss.get_shape()
        loss_num_elements = loss_shape[0].value
        n = math_ops.cast(
            loss_num_elements or array_ops.shape(loss)[0], dtype=loss.dtype)
  
        if have_log_loss:
          log_loss = loss
        else:
          log_loss = math_ops.log(loss)
  
        # Calculate L_hat, Eq. (4) -- stably
        log_mean = math_ops.reduce_logsumexp(log_loss, [0]) - math_ops.log(n)
  
        # expand_dims: Expand shape [a, b, c] to [a, 1, b, c]
        log_loss_expanded = array_ops.expand_dims(log_loss, [1])
  
        # divide: log_loss_sub with shape [a, a, b, c], where
        #
        #  log_loss_sub[i] = log_loss - log_loss[i]
        #
        #       = [ log_loss[j] - log_loss[i] for rows j = 0 ... i - 1     ]
        #         [ zeros                                                  ]
        #         [ log_loss[j] - log_loss[i] for rows j = i + 1 ... a - 1 ]
        #
        log_loss_sub = log_loss - log_loss_expanded
  
        # reduce_sum: Sums each row across all the sub[i]'s; result is:
        #   reduce_sum[j] = (n - 1) * log_loss[j] - (sum_{i != j} loss[i])
        # divide by (n - 1) to get:
        #   geometric_reduction[j] =
        #     log_loss[j] - (sum_{i != j} log_loss[i]) / (n - 1)
        geometric_reduction = math_ops.reduce_sum(log_loss_sub, [0]) / (n - 1)
  
        # subtract this from the original log_loss to get the baseline:
        #   geometric_mean[j] = exp((sum_{i != j} log_loss[i]) / (n - 1))
        log_geometric_mean = log_loss - geometric_reduction
  
        ## Equation (9)
  
        # Calculate sum_{i != j} loss[i] -- via exp(reduce_logsumexp(.))
        # reduce_logsumexp: log-sum-exp each row across all the
        # -sub[i]'s, result is:
        #
        #  exp(reduce_logsumexp[j]) =
        #    1 + sum_{i != j} exp(log_loss[i] - log_loss[j])
        log_local_learning_reduction = math_ops.reduce_logsumexp(
            -log_loss_sub, [0])
  
        # convert local_learning_reduction to the sum-exp of the log-sum-exp
        #  (local_learning_reduction[j] - 1) * exp(log_loss[j])
        #    = sum_{i != j} exp(log_loss[i])
        local_learning_log_sum = (
            _logexpm1(log_local_learning_reduction) + log_loss)
  
        # Add (logaddexp) the local learning signals (Eq. 9)
        local_learning_signal = (
            math_ops.reduce_logsumexp(
                array_ops.stack((local_learning_log_sum, log_geometric_mean)),
                [0])
            - math_ops.log(n))
  
        advantage = log_mean - local_learning_signal
  
        return advantage
  
    return vimco_advantage_fn
function_tokens: def get_vimco_advantage_fn ( have_log_loss = False ) : def vimco_advantage_fn ( _ , loss , name = None ) : """Internal VIMCO function.
  
      Args:
        _: ignored `StochasticTensor`.
        loss: The loss `Tensor`.
        name: Python string, the name scope to use.
  
      Returns:
        The advantage `Tensor`.
      """ with ops . name_scope ( name , "VIMCOAdvantage" , values = [ loss ] ) : loss = ops . convert_to_tensor ( loss ) loss_shape = loss . get_shape ( ) loss_num_elements = loss_shape [ 0 ] . value n = math_ops . cast ( loss_num_elements or array_ops . shape ( loss ) [ 0 ] , dtype = loss . dtype ) if have_log_loss : log_loss = loss else : log_loss = math_ops . log ( loss ) # Calculate L_hat, Eq. (4) -- stably log_mean = math_ops . reduce_logsumexp ( log_loss , [ 0 ] ) - math_ops . log ( n ) # expand_dims: Expand shape [a, b, c] to [a, 1, b, c] log_loss_expanded = array_ops . expand_dims ( log_loss , [ 1 ] ) # divide: log_loss_sub with shape [a, a, b, c], where # #  log_loss_sub[i] = log_loss - log_loss[i] # #       = [ log_loss[j] - log_loss[i] for rows j = 0 ... i - 1     ] #         [ zeros                                                  ] #         [ log_loss[j] - log_loss[i] for rows j = i + 1 ... a - 1 ] # log_loss_sub = log_loss - log_loss_expanded # reduce_sum: Sums each row across all the sub[i]'s; result is: #   reduce_sum[j] = (n - 1) * log_loss[j] - (sum_{i != j} loss[i]) # divide by (n - 1) to get: #   geometric_reduction[j] = #     log_loss[j] - (sum_{i != j} log_loss[i]) / (n - 1) geometric_reduction = math_ops . reduce_sum ( log_loss_sub , [ 0 ] ) / ( n - 1 ) # subtract this from the original log_loss to get the baseline: #   geometric_mean[j] = exp((sum_{i != j} log_loss[i]) / (n - 1)) log_geometric_mean = log_loss - geometric_reduction ## Equation (9) # Calculate sum_{i != j} loss[i] -- via exp(reduce_logsumexp(.)) # reduce_logsumexp: log-sum-exp each row across all the # -sub[i]'s, result is: # #  exp(reduce_logsumexp[j]) = #    1 + sum_{i != j} exp(log_loss[i] - log_loss[j]) log_local_learning_reduction = math_ops . reduce_logsumexp ( - log_loss_sub , [ 0 ] ) # convert local_learning_reduction to the sum-exp of the log-sum-exp #  (local_learning_reduction[j] - 1) * exp(log_loss[j]) #    = sum_{i != j} exp(log_loss[i]) local_learning_log_sum = ( _logexpm1 ( log_local_learning_reduction ) + log_loss ) # Add (logaddexp) the local learning signals (Eq. 9) local_learning_signal = ( math_ops . reduce_logsumexp ( array_ops . stack ( ( local_learning_log_sum , log_geometric_mean ) ) , [ 0 ] ) - math_ops . log ( n ) ) advantage = log_mean - local_learning_signal return advantage return vimco_advantage_fn