Add a torchax container

[yhtang]

No description provided.

[Copilot]

Pull Request Overview

This PR adds support for a TorchAX container to the CI/CD pipeline. TorchAX is being integrated as a new package option alongside existing packages like JAX, XLA, and Flax.

• Adds TorchAX container build job to CI workflow
• Updates package override documentation to include TorchAX
• Creates new Dockerfile for TorchAX container with PyTorch dependencies

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
.github/workflows/ci.yaml	Updates package list documentation and container build matrix to include TorchAX
.github/workflows/_ci.yaml	Adds new build-torchax job with configuration for building TorchAX container
.github/container/manifest.yaml	Adds TorchAX package configuration with GitHub repository and commit tracking
.github/container/Dockerfile.torchax	New Dockerfile defining TorchAX container build with PyTorch and TensorFlow dependencies

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[gpupuck]

Interesting project. Please add some tests (on GPU) in the future