import torch


def bin_data(input, num_bins=None):
    """
    Discretize greyscale values into a finite number of bins
    """
    if num_bins is None:
        return input
    assert num_bins > 0

    # Set each of the corresponding bin indices
    out_data = torch.full_like(input, -1)
    for i in range(num_bins):
        bin_inds = (i / num_bins <= input) * (input <= (i + 1) / num_bins)
        out_data[bin_inds] = i
    assert out_data.max() >= 0

    return out_data.long()