[Test] Fix multi-node fail test and tpu tests

[Michaelvll]

Fix smoke test for multi-node fail fast.
This also fixes a typo for TPU node, that causes test_tpu failure
Also, fixing the issue with the TPU VM tests, due to the API change for jax

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
  • pytest tests/test_smoke.py::test_multi_node_failure
  • pytest tests/test_smoke.py::test_tpu
  • pytest tests/test_smoke.py::test_tpu_vm
  • pytest tests/test_smoke.py::test_tpu_vm_pod
• Backward compatibility tests: bash tests/backward_comaptibility_tests.sh

[romilbhardwaj]

Thanks @Michaelvll! Test passes now!