def sample(self, model_output: torch.Tensor) -> torch.Tensor:
        """
        Sample uniformly between [0, 1] (for each batch example) and return the linear interpolation between the fitted
        quantiles closest to the sampled value.

        model_output is of shape (batch_size, n_timesteps, n_components, n_quantiles)
        """
        device = model_output.device
        num_samples, n_timesteps, n_components, n_quantiles = model_output.shape

        # obtain samples
        probs = torch.rand(
            size=(
                num_samples,
                n_timesteps,
                n_components,
                1,
            )
        ).to(device)
        # add dummy dim
        probas = probs.unsqueeze(-2)

        # tile and transpose
        p = torch.tile(probas, (1, 1, 1, n_quantiles, 1)).transpose(4, 3)

        # prepare quantiles
        tquantiles = torch.tensor(self.quantiles).reshape((1, 1, 1, -1)).to(device)

        # calculate index of biggest quantile smaller than the sampled value
        left_idx = torch.sum(p > tquantiles, dim=-1)

        # obtain index of smallest quantile bigger than sampled value
        right_idx = left_idx + 1

        # repeat the model output on the edges
        repeat_count = [1] * n_quantiles
        repeat_count[0] = 2
        repeat_count[-1] = 2
        repeat_count = torch.tensor(repeat_count).to(device)
        shifted_output = torch.repeat_interleave(model_output, repeat_count, dim=-1)

        # obtain model output values corresponding to the quantiles left and right of the sampled value
        left_value = torch.gather(shifted_output, index=left_idx, dim=-1)
        right_value = torch.gather(shifted_output, index=right_idx, dim=-1)

        # add 0 and 1 to quantiles
        ext_quantiles = [0.0] + self.quantiles + [1.0]
        expanded_q = torch.tile(torch.tensor(ext_quantiles), left_idx.shape).to(device)

        # calculate closest quantiles to the sampled value
        left_q = torch.gather(expanded_q, index=left_idx, dim=-1)
        right_q = torch.gather(expanded_q, index=right_idx, dim=-1)

        # linear interpolation
        weights = (probs - left_q) / (right_q - left_q)
        inter = left_value + weights * (right_value - left_value)

        return inter.squeeze(-1)

        Sample uniformly between [0, 1] (for each batch example) and return the linear interpolation between the fitted
        quantiles closest to the sampled value.

        model_output is of shape (batch_size, n_timesteps, n_components, n_quantiles)
        
darts/utils/likelihood_models/torch.py