message: [tf.data] Handle `tf.SparseTensor` elements in the stats ops.
diff: @@ -20,6 +20,7 @@ from __future__ import print_function
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import iterator_ops
  from tensorflow.python.data.util import nest
  +from tensorflow.python.data.util import sparse
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import gen_dataset_ops
  @@ -161,8 +162,10 @@ class _StatsDataset(dataset_ops.Dataset):
  return self._op_function(
  self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
  self._tag,
  - output_shapes=nest.flatten(self.output_shapes),
  - output_types=nest.flatten(self.output_types))
  + output_types=nest.flatten(
  + sparse.as_dense_types(self.output_types, self.output_classes)),
  + output_shapes=nest.flatten(
  + sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
  @property
  def output_shapes(self):
  