Panic w/ backwards pass when combining gather and max_dim about burn

[nathanielsimard]

I was only able to reproduce the bug on the ndarray backend, it seems to work on the tch backend. You can see the test on the branch: fix/max_dim_gather

[MichaelGoodale]

Hmm, odd that it didn't reproduce with tch! I had whittled down the example to be minimal, and indeed, that one doesn't cause a panic on tch for me either. Perhaps this is two bugs in a trenchcoat pretending to be one!

Here's a specific snippet which does crash for me:

    let a: Vec<f32> = vec![-0.35060948, -0.6759874, -1.2398422, -0.55234957];
    let b = [2, 2, 2, 3, 2, 2, 3, 2, 2, 2, 3, 2, 2, 2, 2, 3, 2, 3, 2, 2];
    let b: Tensor<Autodiff<LibTorch>, 2, Int> =
        Tensor::from_data(Data::from(b.as_slice()), &LibTorchDevice::default()).reshape([5, 4]);
    let a = Tensor::from_data(Data::from(a.as_slice()), &LibTorchDevice::default())
        .reshape([1, 4])
        .require_grad();
    let grammar: Tensor<_, 2> = a.clone().repeat(0, 5);
    let loss = grammar.gather(1, b);
    let loss = loss.clone().max_dim(0) + loss;
    let loss = loss.sum();
    let g = loss.backward();