# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # TODO: b/431034055 - add build-system when OSS build working [project] name = "torch_tpu" version = "0.2.0" license = "Apache-2.0" license-files = ["LICENSE"] # LINT.IfChange(requires_python) requires-python = ">=3.11" # LINT.ThenChange( # bazel/supported_python_versions.bzl:supported_python_versions, # README.md:python_version_support, # ../docs/developers/python_support_policy.md:python_version_support, # ) dependencies = [ "torch>=2.11.0,<=2.15.0", "numpy>=2", "absl-py>=2", # DO NOT change the minimum libtpu version without approval from TorchTPU # leads, as doing so could break partners that are stuck with older # versions. If you have code that needs a newer libtpu version, please make # it conditioned on the version. "libtpu>=0.0.50", "portpicker>=1.6", "tensorboard>=2.20", "frozendict>=2", "immutabledict>=4", "pyzmq>=27", ] [project.optional-dependencies] dev = [ "accelerate>=1.12", "etils[epath]>=1.14", "expecttest>=0.3", "jax >= 0.9.2", "transformers>=5.16.1", ] test = [ "pillow>=11.3.0", "scipy>=1.16.1", "numpy>=2", "timm>=1.0.26", "torchvision>=0.25", "transformers>=5.16.1", "pytest>=9", "six>=1.17", "hypothesis>=6.100", "google-cloud-storage>=2.14.0", "google-api-core>=2.18.0", "unittest-xml-reporting>=4", "cattrs>=24.1.0", ] benchmark = [ "fairscale>=0.4", "flax>=0.12.3", "grain>=0.2.16", "llama-models>=0.3.0", "tensorflow-datasets>=4.9.9", "tensorboardx>=2.6", "timm>=1.0.26", "torchax>=0.0.11", "grpcio-tools>=1.80.0", "dm-tree>=0.1.10", "google-cloud-storage>=2.14.0", "google-api-core>=2.18.0", "transformers>=5.16.1", ] examples = [ "datasets>=4.8", "pillow>=11", "tqdm>=4", "huggingface-hub>=0.34", "transformers>=5.16.1", "diffusers>=0.37", "peft>=0.18", "torchvision>=0.25", "llama-models>=0.3.0", "fairscale>=0.4", "grain>=0.2.16", "tensorflow-datasets>=4.9.9", "accelerate>=1.12", "jax >= 0.9.2", "matplotlib>=3.11.0", "tqdm>=4", "flax>=0.12.3", "tensorboardx>=2.6", "torchax>=0.0.11", "six>=1.17", ] [[tool.uv.index]] name = "pytorch-cpu" url = "https://download.pytorch.org/whl/cpu" explicit = true [tool.uv.sources] torch = { index = "pytorch-cpu" } torchvision = { index = "pytorch-cpu" } [tool.ruff] preview = true force-exclude = true exclude = [ ".git", "build/**", "__pycache__", # Root-level perflab benchmarks (not src/torch_tpu/_internal/benchmarks). "perflab/**", "docs", "examples", ] extend-exclude = [ "**/docs/notebooks/**", "**/docs/html/**", "**/examples/**", ] line-length = 80 indent-width = 2 target-version = "py314" [tool.ruff.lint] ignore = [ # Unnecessary collection call "C408", # Unnecessary map usage "C417", # Unnecessary dict comprehension for iterable "C420", # Object names too complex "C901", # Using `yield` and `return {value}` in a generator function can lead to confusing behavior "B901", # Class could be dataclass or namedtuple "B903", # Raise with from clause inside except block "B904", # Zip without explicit strict parameter "B905", # Google style docstring exceptions: avoid demanding docstrings on __init__ / empty modules (handled by class/package docs) "D100", "D104", "D105", "D107", # PyTorch subclassing exceptions: forward() or dummy helper methods frequently do not use `self` "PLR6301", # Magic numbers (batch sizes, dimensions, thresholds) are frequent in ML code outside strict constant declarations "PLR2004", # TorchTPU packaging structure: submodules and tests import routinely from `torch_tpu._internal` and `torch_tpu._loader` "PLC2701", ] # Selective rule set explicitly chosen to achieve parity with internal `gpylint` and Google Python Style Guide # without triggering noisy false positives on legacy formatting or long URL comments. select = [ # flake8-bugbear opinionated warnings: detects fragile patterns and common logical bugs. "B9", # flake8-comprehensions: catches inefficient or redundant dict/list/set comprehensions. "C", # Pyflakes: checks for undefined variables, syntax errors, and unused imports. Essential for parity with `gpylint`. "F", # pycodestyle warnings: enforces correct whitespace and indentation conventions. "W", # flake8-2020: prevents subtle bugs when checking `sys.version` across 3.11-3.14 targets. "YTT", # flake8-async: catches unsafe concurrency patterns when using asynchronous code. "ASYNC", # Selective pycodestyle errors (`E`): checks indentation without mixing tabs/spaces (`E101`, `E112/113`), # proper comment layout (`E115/117`), and mandatory spaces around operators (`E225/227/228`). "E101", "E112", "E113", "E115", "E117", "E225", "E227", "E228", # Pylint errors (`PLE`): enables real Pylint error detection (undefined variables, signature mismatches) without PLR refactor noise. "PLE", ] [tool.ruff.lint.pydocstyle] convention = "google" [tool.ruff.lint.mccabe] max-complexity = 18 # Pylint configuration (`gpylint` thresholds parity) [tool.ruff.lint.pylint] max-args = 5 max-statements = 50 max-returns = 6 [tool.ruff.lint.per-file-ignores] # Tests and examples rarely require full Google docstrings on every helper class/method/function. "**/tests/**" = ["D", "PLR0913", "PLR2004", "PLR6301"] "**/*_test.py" = ["D", "PLR0913", "PLR2004", "PLR6301"] "**/examples/**" = ["D", "PLR2004"]