tensorflowの計算グラフを描画するwith構文です。このスコープ内のtf.functionなどをtensorboardに描画します。

(tensorflowの解説についての記事を書いている途中なのですが、先にこれだけ載せておきます)

import contextlib
import datetime

import tensorflow as tf

@contextlib.contextmanager
def draw_graph(name: str = 'myfunc'):
    """
    計算グラフをtensorboardに描画する。
    ログは`./tf_function/(name)_YYmmdd-HHMMSS`のフォルダに格納される。
    args:
        name: str(optional)
    """
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    logdir = f'tf_function/{name}_{stamp}'
    writer = tf.summary.create_file_writer(logdir)
    try:
        tf.summary.trace_on()
        yield
    finally:
        with writer.as_default():
            tf.summary.trace_export(name=name, step=0, profiler_outdir=logdir)
            # `tf.summary.trace_off()`は不要(上のtrace_exportで最後にされる)
        writer.close()
        print(f'Saved file to {logdir}.')

Example

with draw_graph('add_and_power_3'):
    @tf.function
    def my_function(x, y):
        """
        f(x, y) = (x + y) ^ 3
        """
        with tf.name_scope('op_add_and_power_3'):
            xpy = x + y
            ret = xpy ** 3
            return ret

    x = tf.constant(1.0)
    y = tf.constant(2.0)
    z = my_function(x, y)
    print('result:',z)

結果:

result: tf.Tensor(27.0, shape=(), dtype=float32)
Saved file to tf_function/add_and_power_3_20211107-191032.