How to mantain the nested structure of a tf.dataset after applying map?
I have managed to solve this problem by &quot;flattening&quot; (eliminating the parenthesis) the return of the random_crop_flip, and applying another map on top of them, where I specified the shapes and returned my desired structure (x ,y), z:
def _set_shapes(img_bl, img_fu, mask):
    img_bl.set_shape([128, 128, 128, 1])
    img_fu.set_shape([128, 128, 128, 1])
    mask.set_shape([128, 128, 128, 1])

    return (img_bl, img_fu), mask 

Then my code looks like this:
x, y = get_filenames(train_data_path, img_type='FLAIR')
z = get_filenames(train_data_path, img_type='mask')

path_dataset = tf.data.Dataset.from_tensor_slices((x, y))
mask_dataset = tf.data.Dataset.from_tensor_slices(z)
dataset = tf.data.Dataset.zip((path_dataset, mask_dataset)).shuffle(50).repeat(10)

ds = dataset. \
    map(lambda xx, zz: ((tf.py_function(load, [xx], [tf.float32, tf.float32])),
                        tf.py_function(load_mask, [zz], [tf.int32])),
        num_parallel_calls=tf.data.AUTOTUNE)

ds = ds.map(lambda xx, zz: (tf.py_function(random_crop_flip, [xx, zz],
                                           [tf.float32, tf.float32, tf.int32])),
            num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.map(_set_shapes)
ds = ds.batch(2)
ds = ds.prefetch(tf.data.AUTOTUNE)