IndexError: tuple index out of range when using a high dimensional Sum_Reduction

[c0g]

I'm using the GeomLoss library and trying to use a Sinkhorn loss with the online backend. For a high dimensional (>=512 I think) tensor it gives me an error when compiling the backward code. Here's the repro with GeomLoss:

from geomloss import SamplesLoss
loss = SamplesLoss(loss="sinkhorn", p=2, blur=.05, backend="online")
x = torch.randn(1000, 512, requires_grad=True)
y = torch.randn(1000, 512)
loss(x, y).backward()

I also have this direct repro by copying values from the debugger:

tagCPUGPU, tag1D2D, tagHostDevice, use_ranges, device_id_request = 1, 0, 0, None, 0
formula, aliases, len_args, dtype = ('Grad_WithSavedForward(Max_SumShiftExp_Reduction(( B - (P * (SqDist(X,Y) / IntCst(2)) ) ),0), Var(0,768,0), Var(4,2,0), Var(5,2,0))', ['X = Vi(768)', 'Y = Vj(768)', 'B = Vj(1)', 'P = Pm(1)', 'Var(4,2,0)', 'Var(5,2,0)'], 6, 'float32')
optional_flags = {'dtype_acc': 'auto', 'sum_scheme': 'block_sum', 'enable_chunks': 1, 'multVar_highdim': 1}

from pykeops.common.keops_io import keops_binder
keops_binder["nvrtc" if tagCPUGPU else "cpp"](tagCPUGPU,
    tag1D2D,
    tagHostDevice,
    use_ranges,
    device_id_request,
    formula,
    aliases,
    len_args,
    dtype,
    "torch",
    optional_flags,
).import_module()

Both give:

File /opt/conda/lib/python3.8/site-packages/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py:103, in GpuReduc1D_finalchunks.get_code(self)
     98 fun_internal = Sum_Reduction(
     99     self.red_formula.formula.children[0], self.red_formula.tagI
    100 )
    101 formula = fun_internal.formula
--> 103 varfinal = self.red_formula.formula.children[1]
    104 nchunks = 1 + (varfinal.dim - 1) // dimfinalchunk
    105 dimlastfinalchunk = varfinal.dim - (nchunks - 1) * dimfinalchunk

IndexError: tuple index out of range

Versions

> python --version
Python 3.8.12
> pip freeze | rg "geom|keops"
geomloss==0.2.5
keopscore==2.1
pykeops==2.1

[jeanfeydy]

Can confirm that the problem now seems to be fixed, so closing the issue. Please feel free to re-open if needed.