ZhiyuLi-goog

Thank you @LeoXinhaoLee for heads up. The warning message looks new to me. I'm not sure how to reproduce this warning message. Would you have any sample code for reproduction?...

Hi @LeoXinhaoLee Many thanks for finding the bug. Have [filed a fix](https://github.com/google/maxtext/pull/495). Feel free to let me know if you have new questions.

The fix was merged in and marked it as fixed.

Closed it with a fix.