import torch


def quantize(batch: torch.Tensor, bins: int) -> torch.Tensor:
    """
    Returns the quantized version of the input batch given a number of bins

    Args:
        batch (torch.Tensor): batch containing data (..., dim_data)
        bins (int): number of intervals for the quantization

    Returns:
        torch.Tensor: quantized tensor containing values from 0 to bins-1
    """

    mins_x = batch.min(dim=-2)[0].unsqueeze(-2)
    maxs_x = batch.max(dim=-2)[0].unsqueeze(-2)

    q = (maxs_x - mins_x) / bins

    index = torch.floor((batch-mins_x) / q)
    index[index==bins] = bins - 1

    return index