See documentation for py_func and eager_py_func.
(func,
inp,
Tout,
stateful=None,
eager=False,
is_grad_func=False,
name=None)
| 263 | |
| 264 | |
| 265 | def _internal_py_func(func, |
| 266 | inp, |
| 267 | Tout, |
| 268 | stateful=None, |
| 269 | eager=False, |
| 270 | is_grad_func=False, |
| 271 | name=None): |
| 272 | """See documentation for py_func and eager_py_func.""" |
| 273 | if not callable(func): |
| 274 | raise ValueError("Expected func to be callable, got func of type {}".format( |
| 275 | type(func))) |
| 276 | |
| 277 | is_list_or_tuple = False |
| 278 | if isinstance(Tout, (list, tuple)): |
| 279 | is_list_or_tuple = True |
| 280 | else: |
| 281 | Tout = [Tout] |
| 282 | |
| 283 | if eager: |
| 284 | func = EagerFunc(func, Tout, is_grad_func) |
| 285 | |
| 286 | token = _py_funcs.insert(func) |
| 287 | # We tie the registered function's lifetime with the current default graph, |
| 288 | # i.e., when the current graph is destroyed, we remove its py funcs. |
| 289 | graph = ops.get_default_graph() |
| 290 | |
| 291 | # pylint: disable=protected-access |
| 292 | while isinstance(graph, function._FuncGraph): |
| 293 | # If the py_func was declared inside a _FuncGraph, its lifetime should be |
| 294 | # bound to that of the outer graph instead. |
| 295 | graph = graph._outer_graph |
| 296 | |
| 297 | # TODO(zhifengc): Consider adding a Graph method to collect |
| 298 | # `cleanup` objects in one of its member. |
| 299 | if not hasattr(graph, "_py_funcs_used_in_graph"): |
| 300 | graph._py_funcs_used_in_graph = [] |
| 301 | |
| 302 | # Store a reference to the function in the graph to ensure it stays alive |
| 303 | # as long as the graph lives. When the graph is destroyed, the function |
| 304 | # is left to the garbage collector for destruction as well. |
| 305 | graph._py_funcs_used_in_graph.append(func) |
| 306 | # pylint: enable=protected-access |
| 307 | |
| 308 | if eager: |
| 309 | result = gen_script_ops.eager_py_func( |
| 310 | input=inp, |
| 311 | token=token, |
| 312 | is_async=context.is_async(), |
| 313 | Tout=Tout, |
| 314 | name=name) |
| 315 | else: |
| 316 | if stateful: |
| 317 | result = gen_script_ops.py_func( |
| 318 | input=inp, token=token, Tout=Tout, name=name) |
| 319 | else: |
| 320 | result = gen_script_ops.py_func_stateless( |
| 321 | input=inp, token=token, Tout=Tout, name=name) |
| 322 | return result if is_list_or_tuple else result[0] |