Retrieve a shared and temporary func graph. The eager execution path lifts a subgraph from the keras global graph into a scratch graph in order to create a function. DistributionStrategies, in turn, constructs multiple functions as well as a final combined function. In order for that logic
(graph=None)
| 499 | |
| 500 | @tf_contextlib.contextmanager |
| 501 | def _scratch_graph(graph=None): |
| 502 | """Retrieve a shared and temporary func graph. |
| 503 | |
| 504 | The eager execution path lifts a subgraph from the keras global graph into |
| 505 | a scratch graph in order to create a function. DistributionStrategies, in |
| 506 | turn, constructs multiple functions as well as a final combined function. In |
| 507 | order for that logic to work correctly, all of the functions need to be |
| 508 | created on the same scratch FuncGraph. |
| 509 | |
| 510 | Args: |
| 511 | graph: A graph to be used as the current scratch graph. If not set then |
| 512 | a scratch graph will either be retrieved or created: |
| 513 | |
| 514 | Yields: |
| 515 | The current scratch graph. |
| 516 | """ |
| 517 | global _CURRENT_SCRATCH_GRAPH |
| 518 | if (_CURRENT_SCRATCH_GRAPH is not None and graph is not None and |
| 519 | _CURRENT_SCRATCH_GRAPH is not graph): |
| 520 | raise ValueError('Multiple scratch graphs specified.') |
| 521 | |
| 522 | if _CURRENT_SCRATCH_GRAPH: |
| 523 | yield _CURRENT_SCRATCH_GRAPH |
| 524 | return |
| 525 | |
| 526 | graph = graph or func_graph.FuncGraph('keras_scratch_graph') |
| 527 | try: |
| 528 | _CURRENT_SCRATCH_GRAPH = graph |
| 529 | yield graph |
| 530 | finally: |
| 531 | _CURRENT_SCRATCH_GRAPH = None |
| 532 | |
| 533 | |
| 534 | @keras_export(v1=['keras.backend.set_session']) |