Tensorflow implementation of Bleu score

I’m looking for a tensorflow implementation of BLEU score
similar to the nltk implementation. The reason I can’t use nltk is
because I need to calculate bleu score per each TPU replica result.
I cannot append predictions across replicas and then use nltk to
calculate BLEU for the entire corpus as I would prefer. The reason
is described in this stackoverflow post
https://stackoverflow.com/questions/60842868/how-can-i-merge-the-results-from-strategy-in-tensorflow-2