In #27679 (comment), I showed how py_function could be extended to handle composite tensor inputs and outputs (and other nested structures, like dicts, tuples, lists, etc). If you have bandwidth to work on a PR that adds that to TensorFlow (with tests etc.), then it would be very welcome; otherwise, you could just use the new_py_function that I defined there, which wraps tf.py_function.

The issue should have been fixed by @edloper. The following example should work.

import tensorflow

# TF1
tf = tensorflow.compat.v1
with tf.Graph().as_default():
    elements = [[1., 2., 3.], [4., 5.]]
    ragged1 = tf.ragged.constant(elements)
    def py_func():
        return tf.ragged.constant(elements)
    
    ragged2 = tf.py_function(
        py_func, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32)
    )
    with tf.Session() as sess:
        print(sess.run(ragged1))  # <tf.RaggedTensorValue [[1.0, 2.0, 3.0], [4.0, 5.0]]>
        print(sess.run(ragged2))  # <tf.RaggedTensorValue [[1.0, 2.0, 3.0], [4.0, 5.0]]>

# TF2
tf = tensorflow.compat.v2
def py_func():
  return tf.ragged.constant([[1., 2., 3.], [4., 5.]])

print(tf.py_function(py_func, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32)))  # <tf.RaggedTensor [[1.0, 2.0, 3.0], [4.0, 5.0]]>

Note the main difference of this example with the example in the original issue is the Tout argument of tf.py_function. When tf.py_function returns a CompositeTensor, the Tout argument should be a subclass of tf.TypeSpec, i.e., tf.RaggedTensorSpec for a RaggedTensor.

See https://www.tensorflow.org/api_docs/python/tf/py_function for more information on tf.py_function and https://www.tensorflow.org/api_docs/python/tf/TypeSpec on tf.TypeSpec.

The issue should have been fixed by @edloper. The following example should work.

import tensorflow

# TF1
tf = tensorflow.compat.v1
with tf.Graph().as_default():
    elements = [[1., 2., 3.], [4., 5.]]
    ragged1 = tf.ragged.constant(elements)
    def py_func():
        return tf.ragged.constant(elements)
    
    ragged2 = tf.py_function(
        py_func, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32)
    )
    with tf.Session() as sess:
        print(sess.run(ragged1))  # <tf.RaggedTensorValue [[1.0, 2.0, 3.0], [4.0, 5.0]]>
        print(sess.run(ragged2))  # <tf.RaggedTensorValue [[1.0, 2.0, 3.0], [4.0, 5.0]]>

# TF2
tf = tensorflow.compat.v2
def py_func():
  return tf.ragged.constant([[1., 2., 3.], [4., 5.]])

print(tf.py_function(py_func, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32)))  # <tf.RaggedTensor [[1.0, 2.0, 3.0], [4.0, 5.0]]>

Note the main difference of this example with the example in the original issue is the Tout argument of tf.py_function. When tf.py_function returns a CompositeTensor, the Tout argument should be a subclass of tf.TypeSpec, i.e., tf.RaggedTensorSpec for a RaggedTensor.

See https://www.tensorflow.org/api_docs/python/tf/py_function for more information on tf.py_function and https://www.tensorflow.org/api_docs/python/tf/TypeSpec on tf.TypeSpec.

Thanks for your brief tutorial, but the given code will not work, if I convert it into the lambda function.

look forward to see the further improvement of tf..

@HuangChiEn Are you using TensorFlow 2.7? Support for using composite tensors (such as RaggedTensor) with py_function was added with 2.7, so if you're using an earlier version of TensorFlow, then it won't work. I tried executing your code as written:

tmp = lambda _: tf.ragged.constant([[1., 2., 3.], [4., 5.]])
tf.py_function(tmp, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32))

And it failed with "<lambda>() missing 1 required positional argument: '_'" (which is expected, since your lambda takes one argument, but you didn't supply any arguments when you called tf.py_function). If I change it the lambda to not expect any argument:

tmp = lambda: tf.ragged.constant([[1., 2., 3.], [4., 5.]])
tf.py_function(tmp, [], Tout=tf.RaggedTensorSpec([2, None], tf.float32))

Then it succeeds for me (in TF 2.7).