pyproject.toml

# Package ######################################################################

[build-system]
# Sync with project.dependencies
requires = ["setuptools", "torch >= 1.13", "numpy", "pybind11 >= 2.10.1"]
build-backend = "setuptools.build_meta"

[project]
name = "torchopt"
description = "An efficient library for differentiable optimization for PyTorch."
readme = "README.md"
# Change this if wheels for `torch` is available
# Search "requires-python" and update all corresponding items
requires-python = ">= 3.7"
authors = [
    { name = "TorchOpt Contributors" },
    { name = "Jie Ren", email = "jieren9806@gmail.com" },
    { name = "Xidong Feng", email = "xidong.feng.20@ucl.ac.uk" },
    { name = "Bo Liu", email = "benjaminliu.eecs@gmail.com" },
    { name = "Xuehai Pan", email = "XuehaiPan@pku.edu.cn" },
]
license = { text = "Apache License, Version 2.0" }
keywords = [
    "PyTorch",
    "functorch",
    "JAX",
    "Meta-Learning",
    "Optimizer",
    "Differentiable Optimizer",
    "Functional Programming",
]
classifiers = [
    "Development Status :: 4 - Beta",
    "License :: OSI Approved :: Apache Software License",
    # Sync with requires-python
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.7",
    "Programming Language :: Python :: 3.8",
    "Programming Language :: Python :: 3.9",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Operating System :: Microsoft :: Windows",
    "Operating System :: POSIX :: Linux",
    "Operating System :: MacOS",
    "Environment :: GPU",
    "Environment :: GPU :: NVIDIA CUDA",
    "Intended Audience :: Developers",
    "Intended Audience :: Education",
    "Intended Audience :: Science/Research",
    "Topic :: Scientific/Engineering",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
    # See also build-system.requires and project.requires-python
    "torch >= 1.13",
    "optree >= 0.4.1",
    "numpy",
    "graphviz",
    "typing-extensions >= 4.0.0",
]
dynamic = ["version"]

[project.urls]
Homepage = "https://github.com/metaopt/torchopt"
Repository = "https://github.com/metaopt/torchopt"
Documentation = "https://torchopt.readthedocs.io"
"Bug Report" = "https://github.com/metaopt/torchopt/issues"

[project.optional-dependencies]
lint = [
    "isort",
    "black[jupyter] >= 22.6.0",
    "pylint[spelling] >= 2.15.0",
    "mypy >= 0.990",
    "types-setuptools",
    "flake8",
    "flake8-bugbear",
    "doc8 < 1.0.0a0",
    "pydocstyle[toml]",
    "pyenchant",
    "cpplint",
    "pre-commit",
]
test = [
    'pytest',
    'pytest-cov',
    'pytest-xdist',
    'jax[cpu] >= 0.3',
    'jaxopt',
    'optax',
]

[tool.setuptools.packages.find]
include = ["torchopt", "torchopt.*"]

# Wheel builder ################################################################
# Reference: https://cibuildwheel.readthedocs.io
[tool.cibuildwheel]
archs = ["auto64"]
build = "*manylinux*"
skip = "pp* *musllinux*"
build-frontend = "pip"
build-verbosity = 3
environment.USE_FP16 = "ON"
environment.CUDACXX = "/usr/local/cuda/bin/nvcc"
environment.TORCH_CUDA_ARCH_LIST = "Common"
environment.DEFAULT_CUDA_VERSION = "11.7"
environment.DEFAULT_TEST_TORCH_SPECS = "cpu cu116"
environment-pass = ["CUDA_VERSION", "TEST_TORCH_SPECS"]
container-engine = "docker"
test-extras = ["test"]

[tool.cibuildwheel.linux]
before-all = """
    CUDA_VERSION="${CUDA_VERSION:-"${DEFAULT_CUDA_VERSION}"}"
    if [[ "${CUDA_VERSION}" == "None" || "${CUDA_VERSION}" == "none" ]]; then
        sed -i -E "s/__version__\\s*=\\s*.*$/\\0 + '+cpu'/" torchopt/version.py
    else
        CUDA_VERSION="$(echo "${CUDA_VERSION}" | cut -d"." -f-2)"
        CUDA_PKG_SUFFIX="$(echo "${CUDA_VERSION}" | tr "." "-")"
        echo "CUDA_VERSION=${CUDA_VERSION}"
        yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo
        yum clean all
        yum install -y nvidia-driver-latest-libs "cuda-minimal-build-${CUDA_PKG_SUFFIX}"
    fi
    echo "cat torchopt/version.py"; cat torchopt/version.py
    touch .first-python
"""
repair-wheel-command = """
    python -m pip install -r requirements.txt
    SITE_PACKAGES="$(python -c 'print(__import__("sysconfig").get_path("purelib"))')"
    TORCH_LIB_PATH="${SITE_PACKAGES}/torch/lib"
    (
        export LD_LIBRARY_PATH="${TORCH_LIB_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/targets/x86_64-linux/lib/stubs${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
        echo "ls ${TORCH_LIB_PATH}"; ls -lh "${TORCH_LIB_PATH}"
        python -m pip install --force-reinstall git+https://github.com/XuehaiPan/auditwheel.git@torchopt
        python -m auditwheel lddtree "{wheel}"
        python -m auditwheel repair --no-copy-site-libs --wheel-dir="{dest_dir}" "{wheel}"
    )
"""
test-command = """
    SITE_PACKAGES="$(python -c 'print(__import__("sysconfig").get_path("purelib"))')"
    TORCH_LIB_PATH="${SITE_PACKAGES}/torch/lib"
    echo "LD_LIBRARY_PATH='${LD_LIBRARY_PATH}'"
    echo "ls ${TORCH_LIB_PATH}"; ls -lh "${TORCH_LIB_PATH}"
    find "${SITE_PACKAGES}/torchopt" -name "*.so" -print0 |
        xargs -0 -I '{}' bash -c "echo 'ldd {}'; ldd '{}'; echo 'patchelf --print-rpath {}'; patchelf --print-rpath '{}'"
    make -C "{project}" test || exit 1
    TORCH_VERSION="$(python -c 'print(__import__("torch").__version__.partition("+")[0])')"
    if [[ -f .first-python ]]; then
        TEST_TORCH_SPECS="${TEST_TORCH_SPECS:-"${DEFAULT_TEST_TORCH_SPECS}"}"
        for spec in ${TEST_TORCH_SPECS}; do
            python -m pip uninstall -y torch
            export PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/${spec}"
            echo "PIP_EXTRA_INDEX_URL='${PIP_EXTRA_INDEX_URL}'"
            python -m pip install "torch==${TORCH_VERSION}"
            echo "ls ${TORCH_LIB_PATH}"; ls -lh "${TORCH_LIB_PATH}"
            find "${SITE_PACKAGES}/torchopt" -name "*.so" -print0 |
                xargs -0 -I '{}' bash -c "echo 'ldd {}'; ldd '{}'; echo 'patchelf --print-rpath {}'; patchelf --print-rpath '{}'"
            make -C "{project}" test || exit 1
        done
        rm -f .first-python
    fi
    rm -rf ~/.pip/cache ~/.cache/pip
"""

# Linter tools #################################################################

[tool.black]
safe = true
line-length = 100
skip-string-normalization = true
# Sync with requires-python
target-version = ["py37", "py38", "py39", "py310", "py311"]

[tool.isort]
atomic = true
profile = "black"
src_paths = ["torchopt", "examples", "tests"]
extra_standard_library = ["typing_extensions"]
indent = 4
line_length = 100
lines_after_imports = 2
multi_line_output = 3

[tool.mypy]
# Sync with requires-python
python_version = 3.7
pretty = true
show_error_codes = true
show_error_context = true
show_traceback = true
allow_redefinition = true
check_untyped_defs = true
disallow_incomplete_defs = false
disallow_untyped_defs = false
ignore_missing_imports = true
no_implicit_optional = true
strict_equality = true
strict_optional = true
warn_no_return = true
warn_redundant_casts = true
warn_unreachable = true
warn_unused_configs = true
warn_unused_ignores = true

[tool.pydocstyle]
convention = "google"

[tool.doc8]
max-line-length = 500