fix: Add JAX Container

[Josephrp]

Linked Issue

#662

Benefits of Having This Container

1. Optimized Performance: JAX uses XLA to compile and run NumPy programs on GPUs, which can significantly speed up numerical computations and machine learning tasks. A container specifically optimized for JAX with CUDA ensures that the environment is configured to leverage GPU acceleration fully.

2. Reproducibility: Containers encapsulate all dependencies, libraries, and configurations needed to run JAX, ensuring that the environment is consistent across different systems. This is crucial for reproducible research and development.

3. Ease of Use: Users can easily pull and run the container without worrying about the complex setup required for GPU support and JAX configuration. This reduces the barrier to entry for new users and accelerates development workflows.

4. Isolation and Security: Containers provide an isolated environment, which enhances security by limiting the impact of potential vulnerabilities. It also avoids conflicts with other software on the host system.

[alexanderankin]

github actions having a day:

1s
55s
Prepare all required actions
Getting action download info
Download action repository 'actions/setup-python@v5' (SHA:39cd[1](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:1)4951b08e74b54015e9e001cdefcf80e669f)
Warning: Failed to download action 'https://api.github.com/repos/actions/setup-python/tarball/39cd14951b08e74b54015e9e001cdefcf80e669f'. Error: Response status code does not indicate success: 504 (Gateway Timeout).
Warning: Back off [2](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:2)4.396 seconds before retry.
Warning: Failed to download action 'https://api.github.com/repos/actions/setup-python/tarball/[3](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:3)9cd14951b08e74b54015e9e001cdefcf80e669f'. Error: Response status code does not indicate success: 502 (Bad Gateway).
Warning: Back off 15.372 seconds before retry.

[Josephrp]

#Issue add Nvidia

Santosh — Today at 04:26
Looks great and thanks for adding this! One recommendation is to use the nightly (ghcr.io/nvidia/jax:jax) from JAX Toolbox and you can file any issues there if tests are failing: https://github.com/NVIDIA/JAX-Toolbox.