class_docstr: A dataset of DTensors.
  
  DTensorDataset encapsulates a `tf.data.Dataset` whose elements are
  automatically packed and returned as DTensors based on a given mesh and
  layouts.
class_docstr_tok: ["A","dataset","of","DTensors",".","DTensorDataset","encapsulates","a","`","tf.data.Dataset","`","whose","elements","are","automatically","packed","and","returned","as","DTensors","based","on","a","given","mesh","and","layouts","."]
docstr: Slices each dataset element on any sharded non-batch dimension.
docstr_tok: ["Slices","each","dataset","element","on","any","sharded","non-batch","dimension","."]
code: def _partition(self, dataset):
      """Slices each dataset element on any sharded non-batch dimension."""
      if (
          self._num_local_devices_per_replica == 1
          and self._partition_offset == 0
      ):
          # Remove this shortcut if tf.data can optimize this away.
          return dataset
  
      # TODO(b/223275517): decouple from self and make testable.
      def slice_batch(index, batch):
          flattened_batch = nest.flatten(batch)
          flattened_output = []
  
          norm_index = math_ops.cast(
              index % self._num_local_devices_per_replica,
              dtype=dtypes.int32,
          )
          norm_index += self._partition_offset
          coords = self._mesh.coords(norm_index)
          coords = array_ops.reshape(coords, (1, -1))
  
          for element, shard_counts, idx_matrix in zip(
              flattened_batch,
              self._all_shard_counts,
              self._index_matrices,
          ):
              indexes = math_ops.matmul(coords, idx_matrix)
              start = array_ops.reshape(indexes, (-1,))
              size = (
                  array_ops.shape_v2(element, out_type=dtypes.int32)
                  // shard_counts
              )
              flattened_output.append(
                  array_ops.slice(element, begin=start, size=size)
              )
  
          return nest.pack_sequence_as(batch, flattened_output)
  
      enumerated_dataset = dataset.enumerate()
      partitioned_dataset = enumerated_dataset.map(slice_batch)
      return partitioned_dataset
  
imports: import dataclasses
  import operator
  from typing import Any, List, Optional, Sequence, Tuple
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import layout
  from tensorflow.python.data.experimental.ops import data_service_ops
  from tensorflow.python.data.experimental.ops import distribute
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import iterator_ops
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.types import data
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 4 for calling the DTensorDataset obj's underlying member method code with example usage: obj._partition(dataset) and returns: partitioned_dataset, dataset, nest