tff.analytics.data_processing.get_capped_elements

Gets the first max_user_contribution words from the input dataset.

The input dataset must yield batched 1-d tensors. This function reads each coordinate of the tensor as an individual element and caps the total number of elements to return. Note either none of the elements in one batch is added to the returned result, or all the elements are added. This means the length of the returned list of elements could be less than max_user_contribution when dataset is capped.

dataset A tf.data.Dataset.
max_user_contribution The maximum number of elements to return.
batch_size The number of elements in each batch of dataset.

A rank-1 Tensor containing the elements of the input dataset after being capped. If the total number of words is less than or equal to max_user_contribution, returns all the words in dataset.