Actually make Python 3.12 work

This PR defaults the dev container to Python 3.12 because Python 3.10 is
soon unusable with future JAX versions.

Python 3.12 builds a package for editing in an isolated pip environment.
That means we need a pyproject.toml file declaring which packages
setup.py depends on, which is added in this PR.

Also setup.py depends on torch to get the C++11 ABI status but that just
won't work in Python 3.12 so that logic is adjusted accordingly. We need
a bunch of other fixes due to the isolated environment change in Python
3.12.

Finally, we need to add `12` into the list of Python versions when
generating JAX dependencies.

TEST: `scripts/build_developer.sh -a -t`

Author: tengyifei

.devcontainer/tpu-contributor/devcontainer.json

@@ -1,12 +1,12 @@
 {
   "name": "tpu-contributor",
-  "image": "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu",
+  "image": "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:3.12_tpuvm",
   "runArgs": [
     "--privileged",
     "--net=host",
     "--shm-size=16G"
   ],
-  "initializeCommand": "docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu",
+  "initializeCommand": "docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:3.12_tpuvm",
   "customizations": {
     "vscode": {
       "extensions": [
@@ -23,4 +23,4 @@
       ]
     }
   }
-}
+}

.devcontainer/tpu-internal/devcontainer.json

@@ -1,16 +1,16 @@
 {
   "name": "tpu-internal",
-  "image": "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu",
+  "image": "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:3.12_tpuvm",
   "runArgs": [
     "--privileged",
     "--net=host",
     "--shm-size=16G"
   ],
   "containerEnv": {
     "BAZEL_REMOTE_CACHE": "1",
-    "SILO_NAME": "cache-silo-${localEnv:USER}-tpuvm"
+    "SILO_NAME": "cache-silo-${localEnv:USER}-tpuvm-312"
   },
-  "initializeCommand": "docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu",
+  "initializeCommand": "docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:3.12_tpuvm",
   "customizations": {
     "vscode": {
       "extensions": [

pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+# These are the packages required to run `setup.py` in an isolated environment.
+# Pip will install these *before* executing the build backend.
+requires = [
+    "setuptools>=42",
+    "wheel",
+    "requests",
+    "numpy",
+    "pyyaml",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-xla"
+description = "XLA bridge for PyTorch"
+readme = "README.md"
+authors = [
+    { name = "PyTorch/XLA Dev Team", email = "[email protected]" },
+]
+license = { file = "LICENSE" }
+requires-python = ">=3.10"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+keywords = ["pytorch", "xla", "tpu", "deep learning", "compiler"]
+
+# This tells build tools to get this info from setup.py instead of this file.
+dynamic = [
+    "version",
+    "dependencies",
+    "optional-dependencies",
+    "entry-points",
+    "scripts"
+]
+
+[project.urls]
+Homepage = "https://github.com/pytorch/xla"
+Repository = "https://github.com/pytorch/xla"
+"Bug Tracker" = "https://github.com/pytorch/xla/issues"

setup.py

@@ -52,6 +52,7 @@
 import contextlib
 import distutils.ccompiler
 import distutils.command.clean
+import importlib.util
 import os
 import re
 import requests
@@ -61,7 +62,13 @@
 import tempfile
 import zipfile
 
-import build_util
+# This gloop imports build_util.py such that it works in Python 3.12's isolated
+# build environment while also not contaminating sys.path which breaks bdist_wheel.
+_PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
+_build_util_path = os.path.join(_PROJECT_DIR, 'build_util.py')
+spec = importlib.util.spec_from_file_location('build_util', _build_util_path)
+build_util = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(build_util)
 
 import platform
 
@@ -151,7 +158,7 @@ def get_git_head_sha(base_dir):
 
 
 def get_build_version(xla_git_sha):
-  version = os.getenv('TORCH_XLA_VERSION', '2.8.0')
+  version = os.getenv('TORCH_XLA_VERSION', '2.9.0')
   if build_util.check_env_flag('GIT_VERSIONED_XLA_BUILD', default='TRUE'):
     try:
       version += '+git' + xla_git_sha[:7]
@@ -270,40 +277,51 @@ def __init__(self, bazel_target):
 class BuildBazelExtension(build_ext.build_ext):
   """A command that runs Bazel to build a C/C++ extension."""
 
-  def run(self):
-    for ext in self.extensions:
-      self.bazel_build(ext)
-    command.build_ext.build_ext.run(self)  # type: ignore
+  def build_extension(self, ext: Extension) -> None:
+    """
+    This method is called by setuptools to build a single extension.
+    We override it to implement our custom Bazel build logic.
+    """
+    if not isinstance(ext, BazelExtension):
+      # If it's not our custom extension type, let setuptools handle it.
+      super().build_extension(ext)
+      return
 
-  def bazel_build(self, ext):
+    # 1. Ensure the temporary build directory exists
     if not os.path.exists(self.build_temp):
       os.makedirs(self.build_temp)
 
+    # 2. Prepare the Bazel command
     bazel_argv = [
         'bazel', 'build', ext.bazel_target,
         f"--symlink_prefix={os.path.join(self.build_temp, 'bazel-')}"
     ]
 
-    build_cpp_tests = build_util.check_env_flag('BUILD_CPP_TESTS', default='0')
-    if build_cpp_tests:
-      bazel_argv.append('//:cpp_tests')
-
-    import torch
-    cxx_abi = os.getenv('CXX_ABI') or getattr(torch._C,
-                                              '_GLIBCXX_USE_CXX11_ABI', None)
-    if cxx_abi is not None:
-      bazel_argv.append(f'--cxxopt=-D_GLIBCXX_USE_CXX11_ABI={int(cxx_abi)}')
+    cxx_abi = os.getenv('CXX_ABI')
+    if cxx_abi is None:
+      try:
+        import torch
+        cxx_abi = getattr(torch._C, '_GLIBCXX_USE_CXX11_ABI', None)
+      except:
+        pass
+    if cxx_abi is None:
+      # Default to building with C++11 ABI, which has been the case since PyTorch 2.7
+      cxx_abi = "1"
+    bazel_argv.append(f'--cxxopt=-D_GLIBCXX_USE_CXX11_ABI={int(cxx_abi)}')
 
     bazel_argv.extend(build_util.bazel_options_from_env())
 
+    # 3. Run the Bazel build
     self.spawn(bazel_argv)
 
+    # 4. Copy the output file to the location setuptools expects
     ext_bazel_bin_path = os.path.join(self.build_temp, 'bazel-bin', ext.relpath,
                                       ext.target_name)
     ext_dest_path = self.get_ext_fullpath(ext.name)
     ext_dest_dir = os.path.dirname(ext_dest_path)
     if not os.path.exists(ext_dest_dir):
       os.makedirs(ext_dest_dir)
+
     shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
 
 
@@ -313,17 +331,21 @@ def bazel_build(self, ext):
   long_description = f.read()
 
 # Finds torch_xla and its subpackages
-packages_to_include = find_packages(include=['torch_xla*'])
-# Explicitly add torchax
-packages_to_include.extend(find_packages(where='torchax', include=['torchax*']))
+# We specify 'torchax.torchax' to find the nested package correctly.
+packages_to_include = find_packages(
+    include=['torch_xla', 'torch_xla.*', 'torchax', 'torchax.*'])
 
 # Map the top-level 'torchax' package name to its source location
-torchax_dir = os.path.join(cwd, 'torchax')
-package_dir_mapping = {'torch_xla': os.path.join(cwd, 'torch_xla')}
-package_dir_mapping['torchax'] = os.path.join(torchax_dir, 'torchax')
+package_dir_mapping = {
+    'torchax': 'torchax/torchax',
+}
 
 
 class Develop(develop.develop):
+  """
+  Custom develop command to build C++ extensions and create a .pth file
+  for a multi-package editable install.
+  """
 
   def run(self):
     # Build the C++ extension
@@ -348,44 +370,42 @@ def link_packages(self):
     (`python setup.py develop`). Nightly and release wheel builds work out of the box
     without egg-link/pth.
     """
+    import glob
+
     # Ensure paths like self.install_dir are set
     self.ensure_finalized()
 
-    # Get the site-packages directory
-    target_dir = self.install_dir
-
-    # Remove the standard .egg-link file
-    # It's usually named based on the distribution name
     dist_name = self.distribution.get_name()
-    egg_link_file = os.path.join(target_dir, dist_name + '.egg-link')
-    if os.path.exists(egg_link_file):
-      print(f"Removing default egg-link file: {egg_link_file}")
-      try:
-        os.remove(egg_link_file)
-      except OSError as e:
-        print(f"Warning: Could not remove {egg_link_file}: {e}")
-
-    # Create our custom .pth file with specific paths
-    cwd = os.path.dirname(__file__)
-    # Path containing 'torch_xla' package source: ROOT
-    path_for_torch_xla = os.path.abspath(cwd)
-    # Path containing 'torchax' package source: ROOT/torchax
-    path_for_torchax = os.path.abspath(os.path.join(cwd, 'torchax'))
-
-    paths_to_add = {path_for_torch_xla, path_for_torchax}
-
-    # Construct a suitable .pth filename (PEP 660 style is good practice)
-    version = self.distribution.get_version()
-    # Sanitize name and version for filename (replace runs of non-alphanumeric chars with '-')
-    sanitized_name = re.sub(r"[^a-zA-Z0-9.]+", "_", dist_name)
-    sanitized_version = re.sub(r"[^a-zA-Z0-9.]+", "_", version)
-    pth_filename = os.path.join(
-        target_dir, f"__editable_{sanitized_name}_{sanitized_version}.pth")
-
-    # Ensure site-packages exists
-    os.makedirs(target_dir, exist_ok=True)
-
-    # Write the paths to the .pth file, one per line
+    install_cmd = self.get_finalized_command('install')
+    target_dir = install_cmd.install_lib
+    assert target_dir is not None
+
+    # Use glob to robustly find and remove the conflicting files.
+    # This is safer than trying to guess the exact sanitized filename.
+    safe_name_part = re.sub(r"[^a-zA-Z0-9]+", "_", dist_name)
+
+    for pattern in [
+        # Remove `.pth` files generated in Python 3.12.
+        f"__editable__.*{safe_name_part}*.pth",
+        f"__editable___*{safe_name_part}*_finder.py",
+        # Also remove the legacy egg-link format.
+        f"{dist_name}.egg-link"
+    ]:
+      for filepath in glob.glob(os.path.join(target_dir, pattern)):
+        print(f"Cleaning up conflicting install file: {filepath}")
+        with contextlib.suppress(OSError):
+          os.remove(filepath)
+
+    # Finally, create our own simple, multi-path .pth file.
+    # We name it simply, e.g., "torch_xla.pth".
+    pth_filename = os.path.join(target_dir, f"{dist_name}.pth")
+
+    project_root = os.path.dirname(os.path.abspath(__file__))
+    paths_to_add = {
+        project_root,  # For `torch_xla`
+        os.path.abspath(os.path.join(project_root, 'torchax')),  # For `torchax`
+    }
+
     with open(pth_filename, "w", encoding='utf-8') as f:
       for path in sorted(paths_to_add):
         f.write(path + "\n")
@@ -403,7 +423,7 @@ def _get_jax_install_requirements():
   jax = f'jax @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jax/jax-{_jax_version}-py3-none-any.whl'
 
   jaxlib = []
-  for python_minor_version in [9, 10, 11]:
+  for python_minor_version in [9, 10, 11, 12]:
     jaxlib.append(
         f'jaxlib @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jaxlib/jaxlib-{_jaxlib_version}-cp3{python_minor_version}-cp3{python_minor_version}-manylinux2014_x86_64.whl ; python_version == "3.{python_minor_version}"'
     )

torch_xla/experimental/gru.py

@@ -99,7 +99,7 @@ def __init__(self, *args, **kwargs):
     super().__init__(*args, **kwargs)
 
   def forward(self, input, hx=None):
-    """
+    r"""
     Args:
         input: Tensor of shape (seq_len, batch, input_size)
         hx: Optional initial hidden state of shape (num_layers, batch, hidden_size).