def cond(pred, true_fn, false_fn):
  """A version of tf.cond that tries to evaluate the condition."""
  v = get_static_value(pred)
  if v is None:
    return tf.cond(pred, true_fn, false_fn)
  if v:
    return true_fn()
  else:
    return false_fn()
A version of tf.cond that tries to evaluate the condition.
trax/tf_numpy/numpy_impl/utils.py