Improve the perf of tf.sparse.segment_(mean/sum/sqrt) ops when using
Improve the perf of tf.sparse.segment_(mean/sum/sqrt) ops when using
bfloat16/float16 precision with float32 for aggregation.

Previously, the memory B/W requirements were as follows for bfloat16/float16
operands using fp32 accumulation:
Input
* Read: #IDs * width * sizeof(bfloat16)
Temporary:
* Read: #Segments * width * sizeof(float32)
* Write: #Segments * width * sizeof(float32)
Output:
* Write: #Segments * width * sizeof(bfloat16)

With the optimization, the memory B/W requirements are as follows:
Input
* Read: #IDs * width * sizeof(bfloat16)
Output:
* Write: #Segments * width * sizeof(bfloat16)

Assuming #IDs = #Segments, and the overhead of writes (exclusive reads + flush)
is 2x that of reads for memory B/W, we reduce memory B/W requirements by 3x.
