pyproject.toml

[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"

[project]
name = "torchprime"
version = "0.1.0"
description = "TorchPrime, reference model implementations for PyTorch on TPU/GPU using torch_xla and torchax"
readme = "README.md"
requires-python = ">=3.10"
license = {file = "LICENSE"}
authors = [
    { name = "PyTorch/XLA team", email = "pytorchxla-dev@google.com" },
]
dependencies = [
    "transformers==4.44.2",
    "transformers[torch]==4.44.2",
    "datasets==3.0.0",
    "hydra-core==1.3.0",
    "optax==0.2.4",
    "fire==0.7.0",
    "tensorflow-cpu==2.18.0",
    "tensorboard==2.18.0",
    "tensorboard-plugin-profile==2.18.0",
    "tf_keras==2.18.0",
]

[project.optional-dependencies]
dev = [
    "ruff~=0.8.6",
    "pytest~=8.3.4",
    "pytest-forked~=1.6.0",
    "click~=8.1.8",
    "toml~=0.10.2",
    "dataclasses-json~=0.6.7",
    "watchdog~=6.0.0",
    "pathspec~=0.12.1",
    "xpk@git+https://github.com/AI-Hypercomputer/xpk@f33d1a6772c3ca73dd68e45b226621534a39b2a5"
]

[project.scripts]
tp = "torchprime.launcher.cli:cli"

[tool.setuptools.packages.find]
where = [""]
include = ["torchprime*"]
exclude = ["torchprime.*.tests.*"]

[tool.setuptools.package-data]
"torchprime" = ["py.typed"]

[tool.pytest.ini_options]
minversion = "6.0"
addopts = "--forked"  # ensure torchax and torch_xla tests don't conflict

[tool.ruff]
indent-width = 2
target-version = "py310"

[tool.ruff.lint]
select = [
    # pycodestyle
    "E",
    # Pyflakes
    "F",
    # pyupgrade
    "UP",
    # flake8-bugbear
    "B",
    # flake8-simplify
    "SIM",
    # isort
    "I",
]
ignore = [
    "E501",  # Line too long. Some copied GPU code has lengthy comments.
]