entmax_bisect bugs with fp16/bf16 about entmax HOT 1 OPEN

As it turns out, we can use the same solution for both torch.float16 and torch.bfloat16. This doesn't seem obvious because they offer different tradeoffs compared to full-precision floats: fp16 keeps the mantissa at the cost of reduced range, bf16 keeps the range at cost of reduced mantissa. But strangely enough we can, using the classic softmax stability trick of subtracting the largest logit from the vector. This makes intuitive sense for torch.float16 because of its reduced range, but what about for torch.bfloat16?

It turns out that bfloat16 has weird problems far away from zero:

x = torch.tensor(0, dtype=torch.bfloat16, device="cuda")
x == (x - 1)  # False, of course

y = torch.tensor(-500, dtype=torch.bfloat16, device="cuda")
y == (y - 1)  # True?!?

So that's why the softmax stability trick works. Bringing it back to our earlier examples:

b = a.to(torch.float16)
entmax.entmax_bisect(b - b.max(), alpha=1.5)  # one-hot!

c = a.to(torch.bfloat16)
entmax.entmax_bisect(c - c.max(), alpha=1.5)  # one-hot!

Tests and PR upcoming.

from entmax.