No graph trace in tensorboard when using @tf.function(jit_compile=True)

[yz-chen18]

Click to expand!

Issue Type

Bug

Source

binary

Tensorflow Version

2.10.0

Custom Code

Yes

OS Platform and Distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

I'm trying to trace a model inference. When I turn off jit_compile option, it works fine. However, when I turn it on, no inference graph appears on tensorboard. I wonder what might happen there. Thanks.

Standalone code to reproduce the issue

logdir = "./xla_dir"

input = np.random.rand(10,3,224,224).astype("float32")

@tf.function(jit_compile=True)
def inf(input):
    # tf_model is a resnet18 model imported from an onnx
    output = tf_model(input_1=input)
    return output

writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True)

inf(input)
with writer.as_default():
    tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

Relevant log output

No response

[bhack]

I am also tracking the same at tensorflow/profiler#503

But the ownership of this issue is still not clear.

/cc @yatbear

[tilakrayal]

@yz-chen18,
As mentioned above, Could you please track the issue tensorflow/profiler#503. Thank you!

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No