Rope optimization part 2

[rwitten]

Single v4-8 step time from 1.209 to 1.199 secs.

(No convergence data since it is actually bit-wise identical locally)

Optimization From Blake!

[rwitten]

The GPU failure is

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: All algorithms tried for %bitcast.220 = bf16[12,2048,8,256]{1,3,2,0} bitcast(bf16[12,8,2048,256]{2,3,1,0} %copy.130), metadata={op_name="jit(train_step)/jit(main)/transpose(jvp(Transformer))/decoder/while/body/remat/decoder/self_attention/query_rotary/convert_element_type[new_dtype=float32 weak_type=False]" source_file="/app/MaxText/layers/embeddings.py" source_line=172} failed. Falling back to default algorithm.  Per-algorithm errors:

when downcasting to bfloat16 on GPU.