Why is the heldout_lp calculated as the mean over all elements of the array?

In the Bayesian neural network example https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/bayesian_neural_network.py, from a theoretical point of view, what does the variable heldout_lp(which I think stands for heldout log probability) represent? Furthermore, why do you take the mean over all elements when calculating heldout_lp?

heldout_lp = np.mean(np.log(mean_probs[np.arange(mean_probs.shape[0]), label_vals.flatten()]))

which could have simply been written as

heldout_lp = np.mean(np.log(mean_probs[:, label_vals.flatten()]))

See line https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/bayesian_neural_network.py#L323.

mean_probs should have the shape (N, 10), where N is the number of test examples, so mean_probs[:, label_vals.flatten()] selects every row of means_probs (so it select all test examples) and, for each of the examples (rows), selects the columns (probabilities) associated with the same test example and all other test examples (label_vals.flatten()). This means that the resulting array of shape (N, N) will have some columns that are equal to each other.

Read tensorflow_probability/examples/bayesian_neural_network.py around line 323 and inspect the shapes of mean_probs and label_vals. Determine whether the indexed expression represents the heldout log probability intended by the example, then document the rationale or correct the example if the indexing is wrong. Confirm the result against the example's test-data dimensions.