fix(training): support training and evaluation on multiple (tpu) workers

[clement-bonnet]

Is your feature request related to a problem? Please describe

Right now, training and evaluators assume local devices only. This works on TPU until a TPU-v-8, but breaks from TPU-v-16 on.

Describe the solution you'd like

Use jax.process_index() to select different keys for each host.

• adapt the evaluator
• adapt the setup_training_state
• adapt the logger to only log for the process_index 0