Skip to content

Remove torchax from torchxla repo. #9672

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 7, 2025
+118 −24,687
Merged

Remove torchax from torchxla repo. #9672

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
38c2ff8
Remove torchax from torchxla repo.
qihqi Oct 6, 2025
61eb0cd
lint
qihqi Oct 6, 2025
95b2a92
new readme
qihqi Oct 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 104 additions & 104 deletions .github/workflows/lintercheck.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,113 +1,113 @@
name: Linter check
on:
pull_request:
push:
branches:
- master
tags:
- r[0-9]+.[0-9]+
pull_request:
push:
branches:
- master
tags:
- r[0-9]+.[0-9]+

jobs:
check_code_changes:
name: Check Code Changes
uses: ./.github/workflows/_check_code_changes.yml
with:
event_name: ${{ github.event_name }}
# For pull_request, use PR's base and head. For push, use event's before and sha.
base_sha: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.event.before }}
head_sha: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
linter_check:
runs-on: ubuntu-24.04
needs: [check_code_changes]
steps:
- name: Checkout repo
if: needs.check_code_changes.outputs.has_code_changes == 'true'
uses: actions/checkout@v3
- name: Setup Python
if: needs.check_code_changes.outputs.has_code_changes == 'true'
uses: actions/setup-python@v4
check_code_changes:
name: Check Code Changes
uses: ./.github/workflows/_check_code_changes.yml
with:
python-version: '3.10'
cache: 'pip'
- run: pip install yapf==0.40.2 # N.B.: keep in sync with `torchax/dev-requirements.txt`, `infra/ansible/config/pip.yaml`
event_name: ${{ github.event_name }}
# For pull_request, use PR's base and head. For push, use event's before and sha.
base_sha: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.event.before }}
head_sha: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
linter_check:
runs-on: ubuntu-24.04
needs: [check_code_changes]
steps:
- name: Checkout repo
if: needs.check_code_changes.outputs.has_code_changes == 'true'
uses: actions/checkout@v3
- name: Setup Python
if: needs.check_code_changes.outputs.has_code_changes == 'true'
uses: actions/setup-python@v4
with:
python-version: "3.10"
cache: "pip"
- run: pip install yapf==0.40.2 # N.B.: keep in sync with `infra/ansible/config/pip.yaml`

- name: Check no TORCH_PIN
if: >
(github.event_name == 'push' && github.event.ref == 'refs/heads/master') &&
needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
run: |
TORCH_PIN=./.torch_pin
if [[ -f "${TORCH_PIN}" ]]; then
echo "Please remove ${TORCH_PIN} before landing."
exit 1
else
echo "No ${TORCH_PIN} found, safe to land..."
fi
- name: Check .cc file extension
shell: bash
run: |
# Find *.cc files recursively in the current directory, limiting to files only.
found_files=$(find . -type f -name "*.cc")
- name: Check no TORCH_PIN
if: >
(github.event_name == 'push' && github.event.ref == 'refs/heads/master') &&
needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
run: |
TORCH_PIN=./.torch_pin
if [[ -f "${TORCH_PIN}" ]]; then
echo "Please remove ${TORCH_PIN} before landing."
exit 1
else
echo "No ${TORCH_PIN} found, safe to land..."
fi
- name: Check .cc file extension
shell: bash
run: |
# Find *.cc files recursively in the current directory, limiting to files only.
found_files=$(find . -type f -name "*.cc")

# Check if any files were found.
if [ -n "$found_files" ]; then
echo "Found *.cc files:"
echo "$found_files"
echo "Please rename them to *.cpp for consistency."
exit 1
else
echo "PASSED *.cc file extension check"
fi
- name: Run clang-format
if: needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
env:
CLANG_FORMAT: clang-format-16
run: |
sudo apt-get update
sudo apt install -y "${CLANG_FORMAT}"
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
echo "Checkout code is not clean"
echo "${git_status}"
exit 1
fi
# Check if any files were found.
if [ -n "$found_files" ]; then
echo "Found *.cc files:"
echo "$found_files"
echo "Please rename them to *.cpp for consistency."
exit 1
else
echo "PASSED *.cc file extension check"
fi
- name: Run clang-format
if: needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
env:
CLANG_FORMAT: clang-format-16
run: |
sudo apt-get update
sudo apt install -y "${CLANG_FORMAT}"
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
echo "Checkout code is not clean"
echo "${git_status}"
exit 1
fi

find . -name '*.cpp' -o -name '*.h' -o -name '*.cc' | xargs "${CLANG_FORMAT}" -i -style=file
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
git diff
echo "${CLANG_FORMAT} recommends the changes above, please manually apply them OR automatically apply the changes "
echo "by running \"${CLANG_FORMAT} -i -style=file /PATH/TO/foo.cpp\" to the following files"
echo "${git_status}"
exit 1
else
echo "PASSED C++ format"
fi
- name: Run yapf
if: needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
run: |
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
echo "Checkout code is not clean"
echo "${git_status}"
exit 1
fi
find . -name '*.cpp' -o -name '*.h' -o -name '*.cc' | xargs "${CLANG_FORMAT}" -i -style=file
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
git diff
echo "${CLANG_FORMAT} recommends the changes above, please manually apply them OR automatically apply the changes "
echo "by running \"${CLANG_FORMAT} -i -style=file /PATH/TO/foo.cpp\" to the following files"
echo "${git_status}"
exit 1
else
echo "PASSED C++ format"
fi
- name: Run yapf
if: needs.check_code_changes.outputs.has_code_changes == 'true'
shell: bash
run: |
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
echo "Checkout code is not clean"
echo "${git_status}"
exit 1
fi

yapf -i -r *.py test/ scripts/ torch_xla/ benchmarks/ torchax/
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
git diff
echo "yapf recommends the changes above, please manually apply them OR automatically apply the changes "
echo "by running `yapf -i /PATH/TO/foo.py` to the following files"
echo "${git_status}"
exit 1
else
echo "PASSED Python format"
fi
- name: Report no code changes
if: needs.check_code_changes.outputs.has_code_changes == 'false'
run: |
echo "No code changes were detected that require running the full test suite."
yapf -i -r *.py test/ scripts/ torch_xla/ benchmarks/
git_status=$(git status --porcelain)
if [[ $git_status ]]; then
git diff
echo "yapf recommends the changes above, please manually apply them OR automatically apply the changes "
echo "by running `yapf -i /PATH/TO/foo.py` to the following files"
echo "${git_status}"
exit 1
else
echo "PASSED Python format"
fi
- name: Report no code changes
if: needs.check_code_changes.outputs.has_code_changes == 'false'
run: |
echo "No code changes were detected that require running the full test suite."
73 changes: 0 additions & 73 deletions .github/workflows/torchax.yml
View file
Open in desktop

This file was deleted.

27 changes: 5 additions & 22 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,8 @@
_jaxlib_version = '0.7.1'
_jax_date = '20250813' # Date for jax and jaxlib.

_torchax_version = '0.0.7' # likely stay the same

if USE_NIGHTLY:
_libtpu_version += f".dev{_libtpu_date}+nightly"
_jax_version += f'.dev{_jax_date}'
Expand Down Expand Up @@ -335,19 +337,6 @@ def build_extension(self, ext: Extension) -> None:
# 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):
"""
Expand All @@ -372,7 +361,7 @@ def link_packages(self):
and `.pth` files. setuptools uses `.egg-link` by default. However, `.egg-link`
only supports linking a single directory containg one editable package.
This function removes the `.egg-link` file and generates a `.pth` file that can
be used to link multiple packages, in particular, `torch_xla` and `torchax`.
be used to link multiple packages.

Note that this function is only relevant in the editable package development path
(`python setup.py develop`). Nightly and release wheel builds work out of the box
Expand Down Expand Up @@ -409,18 +398,13 @@ def link_packages(self):
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")
f.write(project_root + "\n")


def _get_jax_install_requirements():
return [
f'torchax=={_torchax_version}',
f'jaxlib=={_jaxlib_version}',
f'jax=={_jax_version}',
]
Expand Down Expand Up @@ -452,7 +436,6 @@ def _get_jax_install_requirements():
],
python_requires=">=3.10.0",
packages=packages_to_include,
package_dir=package_dir_mapping,
ext_modules=[
BazelExtension('//:_XLAC.so'),
],
Expand Down
28 changes: 0 additions & 28 deletions torchax/LICENSE
View file
Open in desktop

This file was deleted.

Loading