Fix bdist_wheel

Author: tengyifei

scripts/build_developer.sh

@@ -80,6 +80,7 @@ cd $_SCRIPT_DIR/..
 pip uninstall torch_xla torchax torch_xla2 -y
 
 # Build the wheel too, which is useful for other testing purposes.
+rm -f torch_xla.egg-info/SOURCES.txt
 python3 setup.py bdist_wheel
 
 # Link the source files for local development.

setup.py

@@ -335,14 +335,21 @@ def build_extension(self, ext: Extension) -> None:
   long_description = f.read()
 
 # Finds torch_xla and its subpackages
-# 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
-package_dir_mapping = {
-    'torchax': 'torchax/torchax',
-}
+# 1. Find `torch_xla` and its subpackages automatically from the root.
+packages_to_include = find_packages(include=['torch_xla', 'torch_xla.*'])
+
+# 2. Explicitly find the contents of the nested `torchax` package.
+#    Find all sub-packages within the torchax directory (e.g., 'ops').
+torchax_source_dir = 'torchax/torchax'
+torchax_subpackages = find_packages(where=torchax_source_dir)
+#    Construct the full list of packages, starting with the top-level
+#    'torchax' and adding all the discovered sub-packages.
+packages_to_include.extend(['torchax'] +
+                           ['torchax.' + pkg for pkg in torchax_subpackages])
+
+# 3. The package_dir mapping explicitly tells setuptools where the 'torchax'
+#    package's source code begins. `torch_xla` source code is inferred.
+package_dir_mapping = {'torchax': torchax_source_dir}
 
 
 class Develop(develop.develop):