LDA example contains data handling bottleneck.

LDA example, latent_dirichlet_allocation_distributions.py, looks like including a bottleneck in the data loading sequence. By modifying the code like below, the performance was improved. (Runtime log are at the bottom this issue.)
Does that make sense? If necessary, I can send a PR on this. Please feel free to let me know.
(BTW, at least TFP v0.7 (and TF 1.14), this slowdown didn't happen... I'm not sure if this performance change is expected.)

diff --git a/tensorflow_probability/examples/latent_dirichlet_allocation_distributions.py b/tensorflow_probability/examples/latent_dirichlet_allocation_distributions.py
index 2d0c33c5..4f4db7dc 100644
--- a/tensorflow_probability/examples/latent_dirichlet_allocation_distributions.py
+++ b/tensorflow_probability/examples/latent_dirichlet_allocation_distributions.py
@@ -409,7 +409,7 @@ def download(directory, filename):
   return filepath


-def newsgroups_dataset(directory, split_name, num_words, shuffle_and_repeat):
+def newsgroups_dataset(directory, split_name, num_words, batch_size, shuffle_and_repeat):
   """Return 20 newsgroups tf.data.Dataset."""
   data = np.load(download(directory, FILE_TEMPLATE.format(split=split_name)),
                  allow_pickle=True, encoding="latin1")
@@ -440,14 +440,16 @@ def newsgroups_dataset(directory, split_name, num_words, shuffle_and_repeat):
   # stored as a sparse matrix outside of the graph.
   def get_row_py_func(idx):
     def get_row_python(idx_py):
-      return np.squeeze(np.array(sparse_matrix[idx_py].todense()), axis=0)
+      return np.array(sparse_matrix[idx_py].todense())

     py_func = tf.compat.v1.py_func(
         get_row_python, [idx], tf.float32, stateful=False)
-    py_func.set_shape((num_words,))
+    py_func.set_shape((None, num_words))
     return py_func
-  dataset = dataset.map(get_row_py_func)
+  dataset = dataset.batch(batch_size).map(
+    get_row_py_func,
+    num_parallel_calls=tf.data.experimental.AUTOTUNE)
   return dataset


@@ -497,16 +499,17 @@ def build_input_fns(data_dir, batch_size):
   # Build an iterator over training batches.
   def train_input_fn():
     dataset = newsgroups_dataset(
-        data_dir, "train", num_words, shuffle_and_repeat=True)
+        data_dir, "train", num_words, batch_size,
+        shuffle_and_repeat=True)
     # Prefetching makes training about 1.5x faster.
-    dataset = dataset.batch(batch_size).prefetch(32)
+    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
     return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

   # Build an iterator over the heldout set.
   def eval_input_fn():
     dataset = newsgroups_dataset(
-        data_dir, "test", num_words, shuffle_and_repeat=False)
-    dataset = dataset.batch(batch_size)
+        data_dir, "test", num_words, batch_size,
+        shuffle_and_repeat=False)
     return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

   return train_input_fn, eval_input_fn, vocabulary

Runtime log before modification is below.

...(skip)...
INFO:tensorflow:loss = 547.6347, step = 200 (11.816 sec)
I0206 15:08:23.798992 140255794526016 basic_session_run_hooks.py:260] loss = 547.6347, step = 200 (11.816 sec)
INFO:tensorflow:global_step/sec: 8.51187
I0206 15:08:35.536295 140255794526016 basic_session_run_hooks.py:700] global_step/sec: 8.51187
INFO:tensorflow:loss = 632.9415, step = 300 (11.747 sec)
I0206 15:08:35.546317 140255794526016 basic_session_run_hooks.py:260] loss = 632.9415, step = 300 (11.747 sec)
INFO:tensorflow:global_step/sec: 8.53168
I0206 15:08:47.257299 140255794526016 basic_session_run_hooks.py:700] global_step/sec: 8.53168
...(skip)...

Also, runtime log after modification is below.

...(skip)...
INFO:tensorflow:loss = 471.03162, step = 200 (0.661 sec)
I0206 15:13:45.455081 139863353358144 basic_session_run_hooks.py:260] loss = 471.03162, step = 200 (0.661 sec)
INFO:tensorflow:global_step/sec: 146.232
I0206 15:13:46.137598 139863353358144 basic_session_run_hooks.py:700] global_step/sec: 146.232
INFO:tensorflow:loss = 500.88098, step = 300 (0.684 sec)
I0206 15:13:46.139249 139863353358144 basic_session_run_hooks.py:260] loss = 500.88098, step = 300 (0.684 sec)
INFO:tensorflow:global_step/sec: 150.96
I0206 15:13:46.800009 139863353358144 basic_session_run_hooks.py:700] global_step/sec: 150.96
...(skip)...

My environment:

• Docker container: tensorflow/tensorflow:2.1.0-gpu-py3
  • Docker version: 19.03.5
• CPU: Intel(R) Core(TM) i7-6850K CPU @ 3.60GHz
• GPU: RTX 2080 Ti
• Command: python latent_dirichlet_allocation_distributions.py

Start with tensorflow_probability/examples/latent_dirichlet_allocation_distributions.py, especially newsgroups_dataset and build_input_fns. Run the example with the provided command and compare its data-loading and training throughput with the reported logs. Done means determining whether the batching and parallel mapping change addresses the bottleneck without changing the example's behavior.