Computes a distribution over the components for each of the provided datapoints.
- Return type
A tensor of shape
[num_datapoints, num_components]with the assignment probabilities for each component and datapoint. Note that each row of the vector sums to 1, i.e. the returned tensor provides a proper distribution over the components for each datapoint.
When calling this function in a multi-process environment, each process receives only a subset of the predictions. If you want to aggregate predictions, make sure to gather the values returned from this method.