Consider the below nested functions with tf operations and the generated Op Graphs in the following cases:
If the inner function is decorated with
tf.function, it is displayed in the Graph only as generic
If the inner function is not decorated, its operations are expanded and added to the Graph, however without any denoting for being part of a nested process.
Is there a way to display the nested function with its name and then be able to expand its operations when needed? Next I will try to add the functions to a Keras model in order to have a conceptual graph. However, I wonder what are the general possibilities of displaying nested functions with and without decorators and for different types of graphs. Some examples would be really helpful!
logdir = 'logs/func/' + datetime.now().strftime("%Y%m%d-%H%M%S") writer = tf.summary.create_file_writer(logdir) tf.random.set_seed(123) nT = tf.constant(400) n = tf.constant(100000) dt = tf.constant(1/365) drift = tf.constant(0.08) sigma = tf.constant(0.1) K = tf.constant(1.0) @tf.function def brownian_motion(dt, nT, n): dWt = tf.random.normal(mean=0, stddev=tf.math.sqrt(dt), shape=[nT, n]) return dWt @tf.function def option_mc(n, nT, dt, drift, sigma, K): dWt = brownian_motion(dt, nT, n) dYt = drift*dt + sigma*dWt C = tf.cumsum(dYt, axis=0) S = tf.exp(C) A = tf.reduce_mean(S, axis=0) P = tf.reduce_mean(tf.maximum(A - K, 0)) return P tf.summary.trace_on(graph=True, profiler=True) result = option_mc(n, nT, dt, drift, sigma, K) with writer.as_default(): tf.summary.trace_export( name='demo_trace', step=0, profiler_outdir=logdir)