How Cohere uses Ray along with JAX and TPUv4 to train Large Language Models