[Python] Autograph / tf.function produce bad Tensorboard Graph

When using TensorBoard with the following code, the graph generated from tf.function is not fully connected.

I backported the code to tensorflow V1, and it shows the expected graph.

Am I doing something wrong with tf.function? Is there any way to better annotate the operations?


Unrelated Code context:

I'm doing a montecarlo simulation. Given N elements, form K pairs of those elements minimizing the difference. I'm bruteforcing it, testing all possible pairs C(N,2K).


@tf.function
def compute_min_of_max_diff(
samples: tf.Tensor, # [batch, num_candidates]
combs: tf.Tensor, # [num_combinations_chunk, elements]
original_size: tf.Tensor # scalar

) -> tf.Tensor:
sorted_s = tf.sort(samples, axis=1) # [batch, num_candidates]
gathered = tf.gather(sorted_s, combs, axis=1) # [batch, chunk, elements]
diffs = gathered[..., 1::2] - gathered[..., ::2] # [batch, chunk, num_pairs]
max_per_comb = tf.reduce_max(diffs, axis=-1)
return tf.reduce_min(max_per_comb[...,riginal_size], axis=-1)


import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

samples_ph = tf.placeholder(tf.float32)
combs_ph = tf.placeholder(tf.int32)
original_size_ph = tf.placeholder(tf.int32)

sorted_s = tf.sort(samples_ph, axis=1)
gathered = tf.gather(sorted_s, combs_ph, axis=1)
diffs = gathered[..., 1::2] - gathered[..., ::2]
max_per_comb = tf.reduce_max(diffs, axis=-1)
batch_min = tf.reduce_min(max_per_comb[...,riginal_size_ph], axis=-1)
model = batch_min


TensorFlow V2 Graph

Legacy TensorFlow Graph

Continue reading...