Like tf.cond, except emits a single If op.
(pred, true_fn, false_fn, name="cond")
| 56 | |
| 57 | |
| 58 | def cond_v2(pred, true_fn, false_fn, name="cond"): |
| 59 | """Like tf.cond, except emits a single If op.""" |
| 60 | if isinstance(pred, bool): |
| 61 | raise TypeError("pred must not be a Python bool", pred) |
| 62 | |
| 63 | if not name: |
| 64 | name = "cond" |
| 65 | |
| 66 | with ops.name_scope(name) as scope: |
| 67 | true_name = util.unique_fn_name(scope, "true") |
| 68 | false_name = util.unique_fn_name(scope, "false") |
| 69 | |
| 70 | # Automatic control dependencies are added in defuns, but not in v1 |
| 71 | # graphs. Propagate that behavior here. |
| 72 | add_control_dependencies = ops.get_default_graph()._add_control_dependencies |
| 73 | pred = ops.convert_to_tensor(pred) |
| 74 | if (tensor_util.is_tensor(pred) and |
| 75 | (pred.shape.dims is None or pred.shape.dims)): |
| 76 | pred = array_ops.squeeze_v2(pred) |
| 77 | |
| 78 | true_graph = func_graph_module.func_graph_from_py_func( |
| 79 | true_name, |
| 80 | true_fn, [], {}, |
| 81 | func_graph=util.CondBranchFuncGraph( |
| 82 | true_name, collections=ops.get_default_graph()._collections), # pylint: disable=protected-access |
| 83 | add_control_dependencies=add_control_dependencies, |
| 84 | op_return_value=pred) |
| 85 | false_graph = func_graph_module.func_graph_from_py_func( |
| 86 | false_name, |
| 87 | false_fn, [], {}, |
| 88 | func_graph=util.CondBranchFuncGraph( |
| 89 | false_name, collections=ops.get_default_graph()._collections), # pylint: disable=protected-access |
| 90 | add_control_dependencies=add_control_dependencies, |
| 91 | op_return_value=pred) |
| 92 | |
| 93 | verify_captures(_COND, [true_graph, false_graph]) |
| 94 | return _build_cond( |
| 95 | pred, |
| 96 | true_graph, |
| 97 | false_graph, |
| 98 | true_graph.external_captures, |
| 99 | false_graph.external_captures, |
| 100 | building_gradient=False, |
| 101 | name=scope) |
| 102 | |
| 103 | |
| 104 | @ops.RegisterGradient("StatelessIf") |
nothing calls this directly
no test coverage detected