jax:axlearn builds with python 3.12, needs lower python

[dmarx]

Using the system python -- 3.12.3 -- that ships with ghcr.io/nvidia/jax:axlearn-2025-05-16

/opt/axlearn# python -m pip install ".[core,gpu,dev]"
Processing /opt/axlearn
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
Requirement already satisfied: attrs>=23.1.0 in /usr/local/lib/python3.12/dist-packages (25.3.0)
Requirement already satisfied: numpy==1.26.4 in /usr/local/lib/python3.12/dist-packages (1.26.4)
Collecting absl-py==2.1.0
  Using cached absl_py-2.1.0-py3-none-any.whl.metadata (2.3 kB)
Collecting chex==0.1.88
  Using cached chex-0.1.88-py3-none-any.whl.metadata (17 kB)
Collecting importlab==0.8.1
  Using cached importlab-0.8.1-py2.py3-none-any.whl.metadata (2.4 kB)
Collecting jax==0.4.38
  Using cached jax-0.4.38-py3-none-any.whl.metadata (22 kB)
Collecting jaxlib==0.4.38
  Using cached jaxlib-0.4.38-cp312-cp312-manylinux2014_x86_64.whl.metadata (1.0 kB)
Collecting ml-dtypes==0.4.1
  Using cached ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (20 kB)
Requirement already satisfied: msgpack==1.1.0 in /usr/local/lib/python3.12/dist-packages (1.1.0)
Requirement already satisfied: nltk==3.7 in /usr/local/lib/python3.12/dist-packages (3.7)
Collecting optax==0.1.7
  Using cached optax-0.1.7-py3-none-any.whl.metadata (13 kB)
Requirement already satisfied: portpicker in /usr/local/lib/python3.12/dist-packages (1.6.0)
Requirement already satisfied: protobuf>=3.20.3 in /usr/local/lib/python3.12/dist-packages (3.20.3)
Collecting tensorboard-plugin-profile==2.15.1
  Using cached tensorboard_plugin_profile-2.15.1-py3-none-any.whl.metadata (1.0 kB)
Collecting tensorflow==2.17.1
  Using cached tensorflow-2.17.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.2 kB)
Collecting tensorflow-datasets>=4.9.2
  Using cached tensorflow_datasets-4.9.8-py3-none-any.whl.metadata (11 kB)
Collecting tensorflow-io>=0.37.1
  Using cached tensorflow_io-0.37.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (14 kB)
INFO: pip is looking at multiple versions of axlearn[core,dev,gpu] to determine which version is compatible with other requirements. This could take a while.
ERROR: Ignored the following versions that require a different python version: 0.37.0 Requires-Python <3.12,>=3.7
ERROR: Could not find a version that satisfies the requirement tensorflow_text==2.17.0; extra == "core" and platform_machine == "x86_64" (from axlearn[core,dev,gpu]) (from versions: 2.18.1, 2.19.0rc0, 2.19.0)
ERROR: No matching distribution found for tensorflow_text==2.17.0; extra == "core" and platform_machine == "x86_64"

[notice] A new release of pip is available: 23.3.1 -> 25.1.1
[notice] To update, run: python -m pip install --upgrade pip

Not entirely on y'all: the axlearn pyproject.toml should be pinned to whatever version instead of >=3.10 which clearly isn't sufficient.