Failed to convert elements of tf.RaggedTensor in graph building

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text

x_txt = tf.constant(["hello world!"])
bert_pre = hub.load("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
tokenize = hub.KerasLayer(bert_pre.tokenize)
bert_info = bert_pre.tokenize.get_special_tokens_dict()

trimmer = text.RoundRobinTrimmer(max_seq_length=40)
mlm_selector = text.RandomItemSelector(
    max_selections_per_batch=1000,
    selection_rate=0.15,
    unselectable_ids=[
        bert_info["end_of_segment_id"],
        bert_info["mask_id"],
        bert_info["padding_id"],
        bert_info["start_of_sequence_id"],
    ],
)
mask_values_chooser = text.MaskValuesChooser(
    bert_info["vocab_size"], bert_info["mask_id"], mask_token_rate=0.8
)

y_txt = tokenize(x_txt)
y_txt, _segment_id = text.combine_segments(
    [y_txt],
    bert_info["start_of_sequence_id"],
    bert_info["end_of_segment_id"],
)
y_txt = trimmer.trim([y_txt])[0]
print(y_txt)

Running above script (eager execution) results just fine as below

Metal device set to: Apple M1

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

2022-03-23 16:42:31.925269: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-03-23 16:42:31.925539: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2022-03-23 16:42:36.016294: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
2022-03-23 16:42:36.057489: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
2022-03-23 16:42:36.081872: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
2022-03-23 16:42:36.096263: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
<tf.RaggedTensor [[[101],
  [7592],
  [2088],
  [999],
  [102]]]>

However, when I try to compile the graph

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text

x_txt = tf.keras.layers.Input(shape=(), dtype=tf.string)
bert_pre = hub.load("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
tokenize = hub.KerasLayer(bert_pre.tokenize)
bert_info = bert_pre.tokenize.get_special_tokens_dict()

trimmer = text.RoundRobinTrimmer(max_seq_length=40)
mlm_selector = text.RandomItemSelector(
    max_selections_per_batch=1000,
    selection_rate=0.15,
    unselectable_ids=[
        bert_info["end_of_segment_id"],
        bert_info["mask_id"],
        bert_info["padding_id"],
        bert_info["start_of_sequence_id"],
    ],
)
mask_values_chooser = text.MaskValuesChooser(
    bert_info["vocab_size"], bert_info["mask_id"], mask_token_rate=0.8
)

y_txt = tokenize(x_txt)
y_txt, _segment_id = text.combine_segments(
    [y_txt],
    bert_info["start_of_sequence_id"],
    bert_info["end_of_segment_id"],
)
y_txt = trimmer.trim([y_txt])[0]

model = tf.keras.models.Model(inputs=[x_txt], outputs=[y_txt])

Metal device set to: Apple M1

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

2022-03-23 16:43:41.295408: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-03-23 16:43:41.295626: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2022-03-23 16:43:44.494936: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
2022-03-23 16:43:44.530574: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
2022-03-23 16:43:44.554177: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
Traceback (most recent call last):
  File "_text.py", line 31, in <module>
    y_txt = trimmer.trim([y_txt])[0]
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/tensorflow_text/python/ops/trimmer_ops.py", line 49, in trim
    segments = [
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/tensorflow_text/python/ops/trimmer_ops.py", line 50, in <listcomp>
    ragged_tensor.convert_to_tensor_or_ragged_tensor(s) for s in segments
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/tensorflow/python/ops/ragged/ragged_tensor.py", line 2658, in convert_to_tensor_or_ragged_tensor
    return ops.convert_to_tensor_v2_with_dispatch(
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/keras/layers/core/tf_op_layer.py", line 107, in handle
    return TFOpLambda(op)(*args, **kwargs)
  File "/Users/user/miniforge3/envs/tf/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
TypeError: Exception encountered when calling layer "tf.convert_to_tensor" (type TFOpLambda).

Failed to convert elements of tf.RaggedTensor(values=tf.RaggedTensor(values=Tensor("Placeholder:0", shape=(None,), dtype=int32), row_splits=Tensor("Placeholder_1:0", shape=(None,), dtype=int64)), row_splits=Tensor("Placeholder_2:0", shape=(None,), dtype=int64)) to Tensor. Consider casting elements to a supported type. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.

Call arguments received:
  • value=tf.RaggedTensor(values=tf.RaggedTensor(values=Tensor("Placeholder:0", shape=(None,), dtype=int32), row_splits=Tensor("Placeholder_1:0", shape=(None,), dtype=int64)), row_splits=Tensor("Placeholder_2:0", shape=(None,), dtype=int64))
  • dtype=None
  • dtype_hint=None
  • name=None

It outputs this conversion error.

Plus, I'm running this binary for mac m1 but I don't think it is not the cause of the problem.

Start with the supplied graph-building reproduction and inspect tensorflow_text/python/ops/trimmer_ops.py, especially RoundRobinTrimmer.trim and its conversion of segments. Compare the eager and Keras graph paths to identify why the nested tf.RaggedTensor conversion fails; done means constructing the shown model succeeds without the conversion error.