Чередование файлов с subprocess.Popen в tf.data

У меня есть несколько файлов, которые я хотел бы использовать крошечными фрагментами до EOF с tf.data вместо использования tf.read_file один раз для каждого файла (поскольку некоторые файлы намного больше, чем другие).

Я не знаю, как использовать конвейерные подпроцессы в качестве операции TensorFlow (tf.py_func каким-то образом?), А элемент набора данных из list_files известен только во время выполнения графа, поэтому следующее не работает:

def stream(path, bytesize=2048):
    args = f'my_program {path}'
    with subprocess.Popen(args, stdout=subprocess.PIPE) as pipe:
        while True:
            buffer = pipe.stdout.read(bytesize)
            yield np.frombuffer(buffer)
            if len(buffer) < bytesize:
                break


def map_func(path):
    generator = functools.partial(stream, path)
    dataset = tf.data.Dataset.from_generator(generator, tf.float32)
    return dataset


dataset = (
    tf.data.Dataset
    .list_files('data/*')
    .interleave(map_func, batch_size)
    .batch(batch_size)
)

Есть ли способ передать значение элемента набора данных в итерацию, ожидаемую tf.data.Dataset.from_generator, или я поступаю неправильно?

Связанный: Может ли функция карты, предоставленная tf.data.Dataset.from_generator (...) `, разрешить тензорный объект?