Nans gradients for tensors with singleton dimensions for budget_bisect. about entmax HOT 5 OPEN

Hi Saul,

Could you provide some code that causes this error? I haven't been able to reproduce it.

from entmax.

First couple of things I've noticed:

1. The nans occur elsewhere if you change from float64 to float32.
2. The bug does not only occur with budget==2. I can produce it for many inputs with budget==1.

from entmax.

Here is some code:
x = torch.tensor([[[-0.0743, -0.0903]], [[-0.0453, -0.0387]], [[-0.0187, -0.0100]], [[-0.0060, 0.0101]], [[-0.0660, -0.1066]], [[-0.0289, -0.0087]], [[-0.0226, -0.0159]], [[-0.0548, -0.0428]], [[-0.0940, -0.0746]], [[-0.0654, -0.0479]], [[-0.0747, -0.0418]], [[-0.0741, -0.0367]]], requires_grad=True)

check = gradcheck(budget_bisect, (x, 2, -1), eps=1e-5)

from entmax.

Backward also fails if dim == len(x.shape). I believe this problems may be related to the gradient issues.

(it's also something we should have already been checking for, but test_arbitrary_dimension_grad only includes entmax_bisect)

from entmax.