No default summary writer available when using tf.py_function with autograph

[sumanthratna]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): 3.6.8
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS 10.15.5 Beta
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: n/a
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.1.0
• Python version: 3.6.8
• Bazel version (if compiling from source): n/a
• GCC/Compiler version (if compiling from source): n/a
• CUDA/cuDNN version: n/a
• GPU model and memory: n/a

Describe the current behavior Using tf.summary returns False inside a tf.py_function when using autograph.

Describe the expected behavior tf.summary should return True.

Standalone code to reproduce the issue

import tensorflow as tf

@tf.function
def eager_pyfunc():
    def inner_func():
        bool_out = tf.summary.scalar('myscalar', tf.constant(32.9))
        tf.print(bool_out, name='log-myscalar-success')

    tf.py_function(inner_func, [], [], name='log-myscalar')

    bool_out2 = tf.summary.scalar('myscalar2', tf.constant(52.3))
    tf.print(bool_out2, name='log-myscalar2-success')

with tf.summary.create_file_writer('./logs').as_default():
    tf.summary.experimental.set_step(0)
    eager_pyfunc()

outputs:

0
1

Removing the tf.function outputs:

1
1

Other info / logs Could be related to https://github.com/tensorflow/tensorflow/issues/26409.

[ravikyram]

I have tried on colab with TF version 2.1.0, 2.2.0-rc3 and was able to reproduce the issue.Please, find the gist here.Thanks!

[sushreebarsa]

Was able to reproduce the issue in TF 2.14.0-dev20230514,please find the gist here..Thanks !

[Venkat6871]

Hi,

Thank you for opening this issue. Since this issue has been open for a long time, the code/debug information for this issue may not be relevant with the current state of the code base.

The Tensorflow team is constantly improving the framework by fixing bugs and adding new features. We suggest you try the latest TensorFlow version with the latest compatible hardware configuration which could potentially resolve the issue. If you are still facing the issue, please create a new GitHub issue with your latest findings, with all the debugging information which could help us investigate.

Please follow the release notes to stay up to date with the latest developments which are happening in the Tensorflow space.

[sumanthratna]

I re-ran my minimal reproduction sample and this issue still exists on the latest stable TensorFlow.