pyproject.toml

[build-system]
requires = ["hatchling", "hatch-vcs"]
build-backend = "hatchling.build"

[project]
name = "traceax"
dynamic = ["version"]
description = "Stochastic trace estimation in JAX, Lineax, and Equinox"
readme = "README.md"
requires-python = ">=3.9"
license = {text = "Apache2.0"}
keywords = [
  "jax",
  "trace-estimation",
  "statistics",
  "machine-learning",
]
authors = [
  { name = "Linda Serafin", email = "lserafin@usc.edu" },
  { name = "Nicholas Mancuso", email = "nmancuso@usc.edu" },
]
classifiers = [
  "Development Status :: 4 - Beta",
  "Intended Audience :: Science/Research",
  "Intended Audience :: Developers",
  "License :: OSI Approved :: Apache Software License",
  "Natural Language :: English",
  "Topic :: Scientific/Engineering",
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
  "Topic :: Scientific/Engineering :: Information Analysis",
  "Topic :: Scientific/Engineering :: Mathematics",
  "Programming Language :: Python",
  "Programming Language :: Python :: 3.9",
  "Programming Language :: Python :: 3.10",
  "Programming Language :: Python :: 3.11",
  "Programming Language :: Python :: 3.12",
  "Programming Language :: Python :: Implementation :: CPython",
  "Programming Language :: Python :: Implementation :: PyPy",
]
dependencies = [
  "jax>=0.4.13",
  "jaxtyping>=0.2.20",
  "equinox>=0.11.3",
  "lineax>=0.0.4",
  "typing_extensions>=4.5.0",
]

[project.urls]
Documentation = "https://github.com/mancusolab/traceax#readme"
Issues = "https://github.com/mancusolab/traceax/issues"
Source = "https://github.com/mancusolab/traceax"

[tool.hatch.version]
source = "vcs"

[tool.hatch.build.hooks.vcs]
version-file = "src/traceax/_version.py"

[tool.hatch.envs.default]
dependencies = [
  "coverage[toml]>=6.5",
  "pytest",
  "pytest-cov",
]
[tool.hatch.envs.default.scripts]
test = "pytest {args:tests}"
test-cov = "coverage run -m pytest {args:tests}"
cov-report = [
  "- coverage combine",
  "coverage report",
]
cov = [
  "test-cov",
  "cov-report",
]

[[tool.hatch.envs.all.matrix]]
python = ["3.9", "3.10", "3.11", "3.12"]

[tool.hatch.envs.lint]
detached = true
dependencies = [
  "mypy>=1.0.0",
  "ruff>=0.2.2",
]
[tool.hatch.envs.lint.scripts]
typing = "mypy --install-types --non-interactive {args:src/traceax tests}"
style = [
  "ruff {args:.}",
]
fmt = [
  "ruff --fix {args:.}",
  "style",
]
all = [
  "style",
  "typing",
]

[tool.ruff]
line-length = 120

[tool.ruff.lint]
fixable = ["ALL"]
select = ["E", "F", "I001"]
ignore = [
  # Allow non-abstract empty methods in abstract base classes
  "B027",
  # Allow boolean positional values in function calls, like `dict.get(... True)`
  "FBT003",
  # Ignore checks for possible passwords
  "S105", "S106", "S107",
  # Ignore complexity
  "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915",
  # gets confused on empty strings for array types (e.g., Bool[Array, ""])
  "F722"
]
ignore-init-module-imports = true

[tool.ruff.lint.isort]
known-first-party = ["traceax"]
combine-as-imports = true
lines-after-imports = 2
lines-between-types = 1
known-local-folder = ["src"]
section-order = ["future", "standard-library", "third-party", "jax-ecosystem", "first-party", "local-folder"]
extra-standard-library = ["typing_extensions"]
order-by-type = false

[tool.ruff.lint.isort.sections]
jax-ecosystem = ["equinox", "jax", "jaxtyping", "lineax"]

[tool.ruff.lint.per-file-ignores]
# Tests can use magic values, assertions, and relative imports
"tests/**/*" = ["PLR2004", "S101", "TID252"]

[tool.pyright]
reportIncompatibleMethodOverride = true
reportIncompatibleVariableOverride = false  # Incompatible with eqx.AbstractVar
include = ["src/traceax", "tests"]
exclude = ["docs"]

[tool.coverage.run]
source_pkgs = ["traceax", "tests"]
branch = true
parallel = true
omit = []

[tool.coverage.paths]
traceax = ["src/traceax", "*/traceax/src/traceax"]
tests = ["tests", "*/traceax/tests"]

[tool.coverage.report]
exclude_lines = [
  "no cov",
  "if __name__ == .__main__.:",
  "if TYPE_CHECKING:",
]