docstr: Writes a dataset_ops.DatasetV2 to multiple files.
docstr_tok: ["Writes","a","dataset_ops.DatasetV2","to","multiple","files","."]
code: def write_dataset(self, dataset, num_files, num_elems):
      """Writes a dataset_ops.DatasetV2 to multiple files."""
      dataset_paths = []
      dataset_iter = iter(dataset)
  
      for file_idx in range(num_files):
          dataset_path = os.path.join(
              self.get_temp_dir(), f"dataset-{file_idx}.tfrecords"
          )
          dataset_paths.append(dataset_path)
          with tf_record.TFRecordWriter(dataset_path) as writer:
              for _ in range(num_elems // num_files):
                  idx, elem = next(dataset_iter)
                  elem_bytes = example_pb2.Example(
                      features=feature_pb2.Features(
                          feature={
                              "idx": feature_pb2.Feature(
                                  int64_list=feature_pb2.Int64List(
                                      value=[idx]
                                  )
                              ),
                              "elem": feature_pb2.Feature(
                                  bytes_list=feature_pb2.BytesList(
                                      value=[
                                          io_ops.serialize_tensor(
                                              elem
                                          ).numpy()
                                      ]
                                  )
                              ),
                          }
                      )
                  ).SerializeToString()
                  writer.write(elem_bytes)
  
      return dataset_paths
  
imports: import os
  from typing import Any, List, Mapping, Optional, Tuple
  from absl import logging
  from absl.testing import parameterized
  import numpy
  from tensorflow.core.example import example_pb2
  from tensorflow.core.example import feature_pb2
  from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  from tensorflow.dtensor.python.tests import multi_client_test_util
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.experimental.service import server_lib
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import readers
  from tensorflow.python.eager import context
  from tensorflow.python.framework import config
  from tensorflow.python.framework import device_spec
  from tensorflow.python.framework import dtypes
  from tensorflow.python.lib.io import tf_record
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import gen_parsing_ops
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import parsing_config
  from tensorflow.python.ops import parsing_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 3 for calling the MultiClientDTensorDatasetTest obj's underlying member method code with example usage: obj.write_dataset(dataset, num_files, num_elems) and returns: dataset_paths