tensorflow_probability/python/internal/distribution_util.py
def pick_scalar_condition(pred, true_value, false_value, name=None):
  """Convenience function that chooses one of two values based on the predicate.

  This utility is equivalent to a version of `tf.where` that accepts only a
  scalar predicate and computes its result statically when possible. It may also
  be used in place of `tf.cond` when both branches yield a `Tensor` of the same
  shape; the operational difference is that `tf.cond` uses control flow to
  evaluate only the branch that's needed, while `tf.where` (and thus
  this method) may evaluate both branches before the predicate's truth is known.
  This means that `tf.cond` is preferred when one of the branches is expensive
  to evaluate (like performing a large matmul), while this method is preferred
  when both branches are cheap, e.g., constants. In the latter case, we expect
  this method to be substantially faster than `tf.cond` on GPU and to give
  similar performance on CPU.

  Args:
    pred: Scalar `bool` `Tensor` predicate.
    true_value: `Tensor` to return if `pred` is `True`.
    false_value: `Tensor` to return if `pred` is `False`. Must have the same
      shape as `true_value`.
    name: Python `str` name given to ops managed by this object.

  Returns:
    result: a `Tensor` (or `Tensor`-convertible Python value) equal to
      `true_value` if `pred` evaluates to `True` and `false_value` otherwise.
      If the condition can be evaluated statically, the result returned is one
      of the input Python values, with no graph side effects.
  """
  with tf.name_scope(name or "pick_scalar_condition"):
    pred = tf.convert_to_tensor(
        value=pred, dtype_hint=tf.bool, name="pred")
    true_value = tf.convert_to_tensor(value=true_value, name="true_value")
    false_value = tf.convert_to_tensor(value=false_value, name="false_value")
    pred_ = tf.get_static_value(pred)
    if pred_ is None:
      return tf.where(pred, true_value, false_value)
    return true_value if pred_ else false_value
def pick_scalar_condition(pred, true_value, false_value, name=None):
  with tf.name_scope(name or "pick_scalar_condition"):
    pred = tf.convert_to_tensor(
        value=pred, dtype_hint=tf.bool, name="pred")
    true_value = tf.convert_to_tensor(value=true_value, name="true_value")
    false_value = tf.convert_to_tensor(value=false_value, name="false_value")
    pred_ = tf.get_static_value(pred)
    if pred_ is None:
      return tf.where(pred, true_value, false_value)
    return true_value if pred_ else false_value
def pick_scalar_condition ( pred , true_value , false_value , name = None ) : with tf . name_scope ( name or "pick_scalar_condition" ) : pred = tf . convert_to_tensor ( value = pred , dtype_hint = tf . bool , name = "pred" ) true_value = tf . convert_to_tensor ( value = true_value , name = "true_value" ) false_value = tf . convert_to_tensor ( value = false_value , name = "false_value" ) pred_ = tf . get_static_value ( pred ) if pred_ is None : return tf . where ( pred , true_value , false_value ) return true_value if pred_ else false_value
Convenience function that chooses one of two values based on the predicate.

  This utility is equivalent to a version of `tf.where` that accepts only a
  scalar predicate and computes its result statically when possible. It may also
  be used in place of `tf.cond` when both branches yield a `Tensor` of the same
  shape; the operational difference is that `tf.cond` uses control flow to
  evaluate only the branch that's needed, while `tf.where` (and thus
  this method) may evaluate both branches before the predicate's truth is known.
  This means that `tf.cond` is preferred when one of the branches is expensive
  to evaluate (like performing a large matmul), while this method is preferred
  when both branches are cheap, e.g., constants. In the latter case, we expect
  this method to be substantially faster than `tf.cond` on GPU and to give
  similar performance on CPU.

  Args:
    pred: Scalar `bool` `Tensor` predicate.
    true_value: `Tensor` to return if `pred` is `True`.
    false_value: `Tensor` to return if `pred` is `False`. Must have the same
      shape as `true_value`.
    name: Python `str` name given to ops managed by this object.

  Returns:
    result: a `Tensor` (or `Tensor`-convertible Python value) equal to
      `true_value` if `pred` evaluates to `True` and `false_value` otherwise.
      If the condition can be evaluated statically, the result returned is one
      of the input Python values, with no graph side effects.
Convenience function that chooses one of two values based on the predicate .