如何在Tensorflow 2和嵌套列表中使用dataset.map()方法？

import tensorflow as tf

data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
dataset = tf.data.Dataset.from_tensor_slices(data)

def process_element(element):
    # 在这里对元素进行处理
    # 返回处理后的结果
    return element * 2

dataset = dataset.map(process_element)

import tensorflow as tf

data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
dataset = tf.data.Dataset.from_tensor_slices(data)

def process_element(element):
    return element * 2

dataset = dataset.map(process_element)

for element in dataset:
    print(element)

[2, 4, 6]
[8, 10, 12]
[14, 16, 18]