jax

228b61a7 - Add Bzlmod support for JAX.

Add Bzlmod support for JAX. Key changes: - Adds MODULE.bazel to define core dependencies (Abseil, Skylib, Python, XLA/TSL, etc.). - Updates .bazelrc with Bzlmod-specific flags, including --enable_bzlmod and --noenable_workspace. - Configures GitHub Actions to support Bzlmod-based builds and tests. Limitation: - Only support testing against python 3.11 with Bzlmod Imported from GitHub PR: https://github.com/jax-ml/jax/pull/32029 PiperOrigin-RevId: 860049723