Return either `true_fn()` if predicate `pred` is true else `false_fn()`. If `pred` is a bool or has a constant value, we return either `true_fn()` or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both. Arguments: pred: A scalar determining whether to return the result
(pred, true_fn=None, false_fn=None, name=None)
| 25 | |
| 26 | |
| 27 | def smart_cond(pred, true_fn=None, false_fn=None, name=None): |
| 28 | """Return either `true_fn()` if predicate `pred` is true else `false_fn()`. |
| 29 | |
| 30 | If `pred` is a bool or has a constant value, we return either `true_fn()` |
| 31 | or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both. |
| 32 | |
| 33 | Arguments: |
| 34 | pred: A scalar determining whether to return the result of `true_fn` or |
| 35 | `false_fn`. |
| 36 | true_fn: The callable to be performed if pred is true. |
| 37 | false_fn: The callable to be performed if pred is false. |
| 38 | name: Optional name prefix when using `tf.cond`. |
| 39 | |
| 40 | Returns: |
| 41 | Tensors returned by the call to either `true_fn` or `false_fn`. |
| 42 | |
| 43 | Raises: |
| 44 | TypeError: If `true_fn` or `false_fn` is not callable. |
| 45 | """ |
| 46 | if not callable(true_fn): |
| 47 | raise TypeError("`true_fn` must be callable.") |
| 48 | if not callable(false_fn): |
| 49 | raise TypeError("`false_fn` must be callable.") |
| 50 | |
| 51 | pred_value = smart_constant_value(pred) |
| 52 | if pred_value is not None: |
| 53 | if pred_value: |
| 54 | return true_fn() |
| 55 | else: |
| 56 | return false_fn() |
| 57 | else: |
| 58 | return control_flow_ops.cond(pred, true_fn=true_fn, false_fn=false_fn, |
| 59 | name=name) |
| 60 | |
| 61 | |
| 62 | def smart_constant_value(pred): |
nothing calls this directly
no test coverage detected