tf.data.experimental.dense_to_sparse_batch

Like Dataset.padded_batch(), this transformation combines multiple
consecutive elements of the dataset, which might have different
shapes, into a single element. The resulting element has three
components (indices, values, and dense_shape), which
comprise a tf.SparseTensor that represents the same data. The
row_shape represents the dense shape of each row in the
resulting tf.SparseTensor, to which the effective batch size is
prepended. For example:

Args:

batch_size: A tf.int64 scalar tf.Tensor, representing the
number of consecutive elements of this dataset to combine in a
single batch.

row_shape: A tf.TensorShape or tf.int64 vector tensor-like
object representing the equivalent dense shape of a row in the
resulting tf.SparseTensor. Each element of this dataset must
have the same rank as row_shape, and must have size less
than or equal to row_shape in each dimension.