keras/src/backend/tensorflow/core.py/map
def map(f, xs):
    xs = tree.map_structure(convert_to_tensor, xs)

    def get_fn_output_signature(x):
        out = f(x)
        return tree.map_structure(tf.TensorSpec.from_tensor, out)

    fn_output_signature = get_fn_output_signature(xs[0])
    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)