tensorflow: A custom operator get Segmentation Fault in tf.function

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.0.0-beta0-16-g1d91213fe7 2.0.0-beta1
• Python version: 3.6.8
• GCC/Compiler version (if compiling from source): g++ 7.4.0

Describe the current behavior When calling a custom op from a python function with tf.function, I got a segmentation fault. The op run normally without tf.function.

Code to reproduce the issue I implemented a custom op with zero_out_op_kernel_1.cc in tensorflow repository.

I called the operator from the following python code.

import tensorflow as tf

_zero_out_module = tf.load_op_library('custom_ops.so')
zero_out = _zero_out_module.zero_out

@tf.function
def make_zero(x):
    return zero_out(x)

c = tf.constant([4,2,8,9])
res  = make_zero(c)

Other info / logs I confirmed that InferenceContext is NULL. This bug is similar to #30494