tensorflow_probability/python/vi/csiszar_divergence.py
def csiszar_vimco_helper(logu, name=None):
  """Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`.

  `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e.,

  ```none
  logu[j] = log(u[j])
  u[j] = p(x, h[j]) / q(h[j] | x)
  h[j] iid~ q(H | x)
  ```

  Args:
    logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the
      average of `u`. The sum of the gradient of `log_avg_u` is `1`.
    log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the
      average of `u`` except that the average swaps-out `u[i]` for the
      leave-`i`-out Geometric-average. The mean of the gradient of
      `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is,
      ```none
      log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1})
      h[j ; i] = { u[j]                              j!=i
                 { GeometricAverage{u[k] : k != i}   j==i
      ```

  """
  with tf.compat.v1.name_scope(name, "csiszar_vimco_helper", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")

    n = tf.compat.dimension_value(logu.shape.with_rank_at_least(1)[0])
    if n is None:
      n = tf.shape(input=logu)[0]
      log_n = tf.math.log(tf.cast(n, dtype=logu.dtype))
      nm1 = tf.cast(n - 1, dtype=logu.dtype)
    else:
      log_n = np.log(n).astype(logu.dtype.as_numpy_dtype)
      nm1 = np.asarray(n - 1, dtype=logu.dtype.as_numpy_dtype)

    # Throughout we reduce across axis=0 since this is presumed to be iid
    # samples.

    log_max_u = tf.reduce_max(input_tensor=logu, axis=0)
    log_sum_u_minus_log_max_u = tf.reduce_logsumexp(
        input_tensor=logu - log_max_u, axis=0)

    # log_loosum_u[i] =
    # = logsumexp(logu[j] : j != i)
    # = log( exp(logsumexp(logu)) - exp(logu[i]) )
    # = log( exp(logsumexp(logu - logu[i])) exp(logu[i])  - exp(logu[i]))
    # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1)
    # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1)
    # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i])
    d = log_sum_u_minus_log_max_u + (log_max_u - logu)
    # We use `d != 0` rather than `d > 0.` because `d < 0.` should never
    # happens; if it does we want to complain loudly (which `softplus_inverse`
    # will).
    d_ok = tf.not_equal(d, 0.)
    safe_d = tf.where(d_ok, d, tf.ones_like(d))
    d_ok_result = logu + tfd.softplus_inverse(safe_d)

    inf = np.array(np.inf, dtype=logu.dtype.as_numpy_dtype)

    # When not(d_ok) and is_positive_and_largest then we manually compute the
    # log_loosum_u. (We can efficiently do this for any one point but not all,
    # hence we still need the above calculation.) This is good because when
    # this condition is met, we cannot use the above calculation; its -inf.
    # We now compute the log-leave-out-max-sum, replicate it to every
    # point and make sure to select it only when we need to.
    is_positive_and_largest = tf.logical_and(
        logu > 0.,
        tf.equal(logu, log_max_u[tf.newaxis, ...]))
    log_lomsum_u = tf.reduce_logsumexp(
        input_tensor=tf.where(is_positive_and_largest,
                              tf.fill(tf.shape(input=logu), -inf), logu),
        axis=0,
        keepdims=True)
    log_lomsum_u = tf.tile(
        log_lomsum_u,
        multiples=1 + tf.pad(tensor=[n - 1], paddings=[[0, tf.rank(logu) - 1]]))

    d_not_ok_result = tf.where(is_positive_and_largest, log_lomsum_u,
                               tf.fill(tf.shape(input=d), -inf))

    log_loosum_u = tf.where(d_ok, d_ok_result, d_not_ok_result)

    # The swap-one-out-sum ("soosum") is n different sums, each of which
    # replaces the i-th item with the i-th-left-out average, i.e.,
    # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i]))
    #              =  exp(log_loosum_u[i])      + exp(looavg_logu[i])
    looavg_logu = (tf.reduce_sum(input_tensor=logu, axis=0) - logu) / nm1
    log_soosum_u = tf.reduce_logsumexp(
        input_tensor=tf.stack([log_loosum_u, looavg_logu]), axis=0)

    log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n
    log_sooavg_u = log_soosum_u - log_n

    log_avg_u.set_shape(logu.shape.with_rank_at_least(1)[1:])
    log_sooavg_u.set_shape(logu.shape)

    return log_avg_u, log_sooavg_u
def csiszar_vimco_helper(logu, name=None):
  with tf.compat.v1.name_scope(name, "csiszar_vimco_helper", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")

    n = tf.compat.dimension_value(logu.shape.with_rank_at_least(1)[0])
    if n is None:
      n = tf.shape(input=logu)[0]
      log_n = tf.math.log(tf.cast(n, dtype=logu.dtype))
      nm1 = tf.cast(n - 1, dtype=logu.dtype)
    else:
      log_n = np.log(n).astype(logu.dtype.as_numpy_dtype)
      nm1 = np.asarray(n - 1, dtype=logu.dtype.as_numpy_dtype)

    # Throughout we reduce across axis=0 since this is presumed to be iid
    # samples.

    log_max_u = tf.reduce_max(input_tensor=logu, axis=0)
    log_sum_u_minus_log_max_u = tf.reduce_logsumexp(
        input_tensor=logu - log_max_u, axis=0)

    # log_loosum_u[i] =
    # = logsumexp(logu[j] : j != i)
    # = log( exp(logsumexp(logu)) - exp(logu[i]) )
    # = log( exp(logsumexp(logu - logu[i])) exp(logu[i])  - exp(logu[i]))
    # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1)
    # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1)
    # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i])
    d = log_sum_u_minus_log_max_u + (log_max_u - logu)
    # We use `d != 0` rather than `d > 0.` because `d < 0.` should never
    # happens; if it does we want to complain loudly (which `softplus_inverse`
    # will).
    d_ok = tf.not_equal(d, 0.)
    safe_d = tf.where(d_ok, d, tf.ones_like(d))
    d_ok_result = logu + tfd.softplus_inverse(safe_d)

    inf = np.array(np.inf, dtype=logu.dtype.as_numpy_dtype)

    # When not(d_ok) and is_positive_and_largest then we manually compute the
    # log_loosum_u. (We can efficiently do this for any one point but not all,
    # hence we still need the above calculation.) This is good because when
    # this condition is met, we cannot use the above calculation; its -inf.
    # We now compute the log-leave-out-max-sum, replicate it to every
    # point and make sure to select it only when we need to.
    is_positive_and_largest = tf.logical_and(
        logu > 0.,
        tf.equal(logu, log_max_u[tf.newaxis, ...]))
    log_lomsum_u = tf.reduce_logsumexp(
        input_tensor=tf.where(is_positive_and_largest,
                              tf.fill(tf.shape(input=logu), -inf), logu),
        axis=0,
        keepdims=True)
    log_lomsum_u = tf.tile(
        log_lomsum_u,
        multiples=1 + tf.pad(tensor=[n - 1], paddings=[[0, tf.rank(logu) - 1]]))

    d_not_ok_result = tf.where(is_positive_and_largest, log_lomsum_u,
                               tf.fill(tf.shape(input=d), -inf))

    log_loosum_u = tf.where(d_ok, d_ok_result, d_not_ok_result)

    # The swap-one-out-sum ("soosum") is n different sums, each of which
    # replaces the i-th item with the i-th-left-out average, i.e.,
    # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i]))
    #              =  exp(log_loosum_u[i])      + exp(looavg_logu[i])
    looavg_logu = (tf.reduce_sum(input_tensor=logu, axis=0) - logu) / nm1
    log_soosum_u = tf.reduce_logsumexp(
        input_tensor=tf.stack([log_loosum_u, looavg_logu]), axis=0)

    log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n
    log_sooavg_u = log_soosum_u - log_n

    log_avg_u.set_shape(logu.shape.with_rank_at_least(1)[1:])
    log_sooavg_u.set_shape(logu.shape)

    return log_avg_u, log_sooavg_u
def csiszar_vimco_helper ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "csiszar_vimco_helper" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) n = tf . compat . dimension_value ( logu . shape . with_rank_at_least ( 1 ) [ 0 ] ) if n is None : n = tf . shape ( input = logu ) [ 0 ] log_n = tf . math . log ( tf . cast ( n , dtype = logu . dtype ) ) nm1 = tf . cast ( n - 1 , dtype = logu . dtype ) else : log_n = np . log ( n ) . astype ( logu . dtype . as_numpy_dtype ) nm1 = np . asarray ( n - 1 , dtype = logu . dtype . as_numpy_dtype ) # Throughout we reduce across axis=0 since this is presumed to be iid # samples. log_max_u = tf . reduce_max ( input_tensor = logu , axis = 0 ) log_sum_u_minus_log_max_u = tf . reduce_logsumexp ( input_tensor = logu - log_max_u , axis = 0 ) # log_loosum_u[i] = # = logsumexp(logu[j] : j != i) # = log( exp(logsumexp(logu)) - exp(logu[i]) ) # = log( exp(logsumexp(logu - logu[i])) exp(logu[i])  - exp(logu[i])) # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1) # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1) # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i]) d = log_sum_u_minus_log_max_u + ( log_max_u - logu ) # We use `d != 0` rather than `d > 0.` because `d < 0.` should never # happens; if it does we want to complain loudly (which `softplus_inverse` # will). d_ok = tf . not_equal ( d , 0. ) safe_d = tf . where ( d_ok , d , tf . ones_like ( d ) ) d_ok_result = logu + tfd . softplus_inverse ( safe_d ) inf = np . array ( np . inf , dtype = logu . dtype . as_numpy_dtype ) # When not(d_ok) and is_positive_and_largest then we manually compute the # log_loosum_u. (We can efficiently do this for any one point but not all, # hence we still need the above calculation.) This is good because when # this condition is met, we cannot use the above calculation; its -inf. # We now compute the log-leave-out-max-sum, replicate it to every # point and make sure to select it only when we need to. is_positive_and_largest = tf . logical_and ( logu > 0. , tf . equal ( logu , log_max_u [ tf . newaxis , ... ] ) ) log_lomsum_u = tf . reduce_logsumexp ( input_tensor = tf . where ( is_positive_and_largest , tf . fill ( tf . shape ( input = logu ) , - inf ) , logu ) , axis = 0 , keepdims = True ) log_lomsum_u = tf . tile ( log_lomsum_u , multiples = 1 + tf . pad ( tensor = [ n - 1 ] , paddings = [ [ 0 , tf . rank ( logu ) - 1 ] ] ) ) d_not_ok_result = tf . where ( is_positive_and_largest , log_lomsum_u , tf . fill ( tf . shape ( input = d ) , - inf ) ) log_loosum_u = tf . where ( d_ok , d_ok_result , d_not_ok_result ) # The swap-one-out-sum ("soosum") is n different sums, each of which # replaces the i-th item with the i-th-left-out average, i.e., # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i])) #              =  exp(log_loosum_u[i])      + exp(looavg_logu[i]) looavg_logu = ( tf . reduce_sum ( input_tensor = logu , axis = 0 ) - logu ) / nm1 log_soosum_u = tf . reduce_logsumexp ( input_tensor = tf . stack ( [ log_loosum_u , looavg_logu ] ) , axis = 0 ) log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n log_sooavg_u = log_soosum_u - log_n log_avg_u . set_shape ( logu . shape . with_rank_at_least ( 1 ) [ 1 : ] ) log_sooavg_u . set_shape ( logu . shape ) return log_avg_u , log_sooavg_u
Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`.

  `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e.,

  ```none
  logu[j] = log(u[j])
  u[j] = p(x, h[j]) / q(h[j] | x)
  h[j] iid~ q(H | x)
  ```

  Args:
    logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the
      average of `u`. The sum of the gradient of `log_avg_u` is `1`.
    log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the
      average of `u`` except that the average swaps-out `u[i]` for the
      leave-`i`-out Geometric-average. The mean of the gradient of
      `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is,
      ```none
      log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1})
      h[j ; i] = { u[j]                              j!=i
                 { GeometricAverage{u[k] : k != i}   j==i
      ```
Helper to csiszar_vimco ; computes log_avg_u log_sooavg_u .