Commit 2cd247d2 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

Handle positive and negative infinity in TopKV2.

TopKV2 hides iota in the low bits of the input after converting from bf16 to f32. This usually works, but for positive and negative infinity or'ing in iota produces NANs.

To handle positive and negative infinity, treat bf16 as integers in
sign-magnitude format. Convert to two's complement. Sort in two's complement and
convert back.

Add an exhaustive unit test for bfloat16 to float conversion.

PiperOrigin-RevId: 201421784
parent 164099ee
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment