tfp.experimental.distributions.marginal_fns.ps.smart_where

As tf.where, but only calls x_fn/y_fn when condition not statically known.

tfp.experimental.distributions.marginal_fns.ps.smart_where(
    condition, x_fn, y_fn
)

Important: Since this avoids executing the inoperative branch when possible, it will not necessarily broadcast x_fn() with y_fn(), so it is imperative that they return Tensors which broadcast with condition to the same final shape.

Args
condition	A bool Tensor.
x_fn	A callable returning a Tensor, for locations where condition is True.
y_fn	A callable returning a Tensor, for locations where condition is False.

Returns
A Tensor equivalent to tf.where(condition, x_fn(), y_fn()).