No accelerators in lr calculation in Trax step (to avoid allocations).
No accelerators in lr calculation in Trax step (to avoid allocations).

PiperOrigin-RevId: 259968721

tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor