Array index out of range in "rational_quadratic" file

[Yifei-Xiong]

Hi,
I'm having some issues with this package, in file: "nflows/transforms/splines/rational_quadratic.py"
Line 114: input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
the index bin_idx may out of range. It takes values from {0, 1, ..., num_bins}, but the shape of widths is batch * num_bins

I found a possible reason, in Line 78 and Line 79:

if torch.min(inputs) < left or torch.max(inputs) > right:
    raise InputOutsideDomain()

if some values in inputs equal to right, it also satisfied this range check, but this point is actually at the right endpoint of the last bin, that is, the bin_idx generated in Line 111
bin_idx = torchutils.searchsorted(cumwidths, inputs)[..., None]
will give the corresponding bin_idx value as num_bins
(In this case, if I set inputs -= 1e-5, the corresponding bin_idx will be num_bins-1)
Now, if bin_idx has some value equals to num_bins, then in Line 114 input_bin_widths = widths.gather(-1, bin_idx)[..., 0] can lead to an exception because the shape of widths is batch * num_bins

I think one possible solution is that, add bin_idx[bin_idx == num_bins] -= 1 after setting the values of bin_idx, that is

if inverse:
    bin_idx = torchutils.searchsorted(cumheights, inputs)[..., None]
else:
    bin_idx = torchutils.searchsorted(cumwidths, inputs)[..., None]

bin_idx[bin_idx == num_bins] -= 1
input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
input_bin_widths = widths.gather(-1, bin_idx)[..., 0]