src/transformers/models/esm/modeling_esmfold.py/EsmCategoricalMixture/__init__
class EsmCategoricalMixture:    def __init__(self, param, bins=50, start=0, end=1):
        # All tensors are of shape ..., bins.
        self.logits = param
        bins = torch.linspace(start, end, bins + 1, device=self.logits.device, dtype=self.logits.dtype)
        self.v_bins = (bins[:-1] + bins[1:]) / 2