Maxtext unit tests with Pathways backend.

[RoshaniN]

Description

This PR enables TPU unit tests to also run with Pathways backend. Essentially, we will have two sets of tests - one with McJAX and one with Pathways.

• This change being made to ensure feature parity between Pathways and McJAX.
• The tests run as part of a docker compose script which sets up the Pathways containers along with Maxtext. (Github Actions didn't have enough support for deploying Pathways containers as "service containers".)
• TPU integration tests may be also run with Pathways backend in the future.

For more details, please read the doc on b/397475777 . Note that extra self-hosted runners have been added so that tests can be executed in parallel and complete faster overall.

Tests

Please describe how you tested this change -

1. Changes tested locally using command bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
2. Pathways flow tested on Github workflow - Example runs -

• https://github.com/AI-Hypercomputer/maxtext/actions/runs/13444788084/job/37567625491?pr=1211 (pathways only)
• https://github.com/AI-Hypercomputer/maxtext/actions/runs/13557045534/job/37893343616 (entire workflow)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• [X] I have performed a self-review of my code.
• [X] I have necessary comments in my code, particularly in hard-to-understand areas.
• [X] I have run end-to-end tests tests and provided workload links above if applicable.
• [X] I have made or will make corresponding changes to the doc if needed.