pyproject.toml

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "tjax"
version = "1.1.3"
description = "Tools for JAX."
readme = "README.rst"
requires-python = ">=3.11,<3.14"
license = {file = "LICENSE"}
authors = [{email = "mistersheik@gmail.com"}, {name = "Neil Girdhar"}]
maintainers = [{email = "mistersheik@gmail.com"}, {name = "Neil Girdhar"}]
classifiers = [
  "Development Status :: 4 - Beta",
  "Intended Audience :: Science/Research",
  "License :: OSI Approved :: Apache Software License",
  "Operating System :: OS Independent",
  "Programming Language :: Python :: 3",
  "Programming Language :: Python :: 3.11",
  "Programming Language :: Python :: 3.12",
  "Programming Language :: Python :: 3.13",
  "Programming Language :: Python",
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
  "Topic :: Software Development :: Libraries :: Python Modules",
  "Topic :: Utilities",
  "Typing :: Typed",
]
dependencies = [
  "array-api-extra>=0.6.0",
  "array_api_compat>=1.10",
  "jax>=0.4.27",
  "numpy>=1.25",
  "optax>=0.2.3",
  "rich>=13.7",
  "typing_extensions>=4.11",
]

[project.urls]
repository = "https://github.com/NeilGirdhar/tjax"

[tool.uv]
dev-dependencies = [
  "chex>=0.1.3",
  "flax>=0.10.3",
  "isort>=5.13",
  "jupyter>=1",
  "mypy>=1.12",
  "torch>=2.6",
  "networkx>=3.1",
  "pre-commit>=4",
  "pylint>=3.3",
  "pyright>=0.0.13",
  "pytest>=8",
  "ruff>=0.8",
]

[tool.isort]
line_length = 100
multi_line_output = 0
skip_glob = [".pytype/*.py"]

[tool.pylint.master]
jobs = 0
py-version = "3.11"
load-plugins = [
  "pylint.extensions.bad_builtin",
  "pylint.extensions.broad_try_clause",
  "pylint.extensions.check_elif",
  "pylint.extensions.code_style",
  "pylint.extensions.comparison_placement",
  "pylint.extensions.confusing_elif",
  "pylint.extensions.consider_refactoring_into_while_condition",
  "pylint.extensions.consider_ternary_expression",
  "pylint.extensions.dict_init_mutate",
  "pylint.extensions.docparams",
  "pylint.extensions.docstyle",
  "pylint.extensions.dunder",
  "pylint.extensions.empty_comment",
  "pylint.extensions.eq_without_hash",
  "pylint.extensions.for_any_all",
  "pylint.extensions.magic_value",
  "pylint.extensions.mccabe",
  "pylint.extensions.no_self_use",
  "pylint.extensions.overlapping_exceptions",
  "pylint.extensions.private_import",
  "pylint.extensions.redefined_loop_name",
  "pylint.extensions.redefined_variable_type",
  "pylint.extensions.set_membership",
  "pylint.extensions.typing",
  "pylint.extensions.while_used",
]

[tool.pylint.messages_control]
disable = [
  "abstract-method",
  "arguments-differ",
  "bad-dunder-name",
  "confusing-consecutive-elif",
  "consider-using-assignment-expr",
  "consider-using-tuple",
  "cyclic-import",
  "duplicate-code",
  "empty-comment",
  "fixme",
  "import-outside-toplevel",
  "invalid-field-call",
  "invalid-name",
  "invalid-unary-operand-type",
  "line-too-long",
  "keyword-arg-before-vararg",
  "magic-value-comparison",
  "missing-docstring",
  "no-member",
  "no-self-use",
  "not-callable",
  "protected-access",
  "redefined-builtin",
  "redefined-outer-name",
  "redefined-variable-type",
  "signature-differs",
  "too-complex",
  "too-few-public-methods",
  "too-many-ancestors",
  "too-many-arguments",
  "too-many-branches",
  "too-many-instance-attributes",
  "too-many-locals",
  "too-many-positional-arguments",
  "too-many-return-statements",
  "too-many-statements",
  "too-many-try-statements",
  "unexpected-keyword-arg",
  "unused-argument",
  "unused-import",
  "use-set-for-membership",
  "while-used",
]
enable = [
  "useless-suppression",
  "use-symbolic-message-instead",
]

[tool.pylint.reports]
output-format = "colorized"

[tool.pylint.similarities]
min-similarity-lines = 10
ignore-imports = "yes"

[tool.pyright]
include = ["tjax", "tests"]
stubPath = ""
venvPath = "."
venv = ".venv"
enableTypeIgnoreComments = false
reportImportCycles = true
reportCallInDefaultInitializer = true
reportConstantRedefinition = true
reportDeprecated = true
reportDuplicateImport = true
reportImplicitOverride = true
reportImplicitStringConcatenation = false
reportIncompatibleMethodOverride = true
reportIncompatibleVariableOverride = true
reportInconsistentConstructor = true
reportMatchNotExhaustive = true
reportMissingParameterType = true
reportMissingSuperCall = true
reportMissingTypeArgument = true
reportOverlappingOverload = true
reportPrivateImportUsage = true
reportPropertyTypeMismatch = true
reportShadowedImports = true
reportUninitializedInstanceVariable = true
reportUnknownArgumentType = false
reportUnknownLambdaType = false
reportUnknownMemberType = false
reportUnknownParameterType = false
reportUnknownVariableType = false
reportUnnecessaryCast = true
reportUnnecessaryComparison = true
reportUnnecessaryContains = true
reportUnnecessaryIsInstance = true
reportUnnecessaryTypeIgnoreComment = true
reportUntypedBaseClass = true
reportUntypedClassDecorator = true
reportUntypedFunctionDecorator = true
reportUntypedNamedTuple = true
reportUnusedCallResult = false
reportUnusedClass = true
reportUnusedExpression = true
reportUnusedFunction = true
reportUnusedVariable = true

[tool.mypy]
files = ["tjax", "tests"]
check_untyped_defs = true
disallow_any_generics = true
disallow_incomplete_defs = true
disallow_untyped_calls = true
disallow_untyped_decorators = true
disallow_untyped_defs = true
no_implicit_optional = true
pretty = true
show_error_codes = true
show_error_context = false
strict_equality = true
warn_no_return = true
warn_redundant_casts = true
warn_return_any = true
warn_unreachable = true
warn_unused_configs = true
warn_unused_ignores = true

[[tool.mypy.overrides]]
module = ["networkx", "optax", "optax.contrib", "array_api_compat"]
ignore_missing_imports = true

[tool.ruff]
line-length = 100

[tool.ruff.lint]
preview = true
select = ["ALL"]
ignore = [
  "ANN401", # Dynamically typed expressions (Any).
  "ARG001", # Unused function argument.
  "ARG002", # Unused method argument.
  "ARG003", # Unused class method argument.
  "ARG004", # Unused static method argument.
  "B011", # Do not assert false.
  "C901", # Complex structure.
  "COM812", # Trailing comma missing.
  "CPY001", # Missing copyright.
  "D100", # Missing docstring in public module.
  "D101", # Missing docstring in public class.
  "D102", # Missing docstring in public method.
  "D103", # Missing docstring in public function.
  "D105", # Missing docstring in magic method.
  "D107", # Missing docstring in __init__.
  "DOC201", # return is not documented in docstring.
  "DOC202", # Docstring should not have a returns section because the function doesn"t return anything.
  "DOC402", # DOC402 `yield` is not documented in docstring.
  "DOC501", # Raised exception `ValueError` missing from docstring.
  "ERA001", # Commented-out code.
  "F722", # Syntax error in forward annotation.
  "FBT003", # Boolean positional value in function call.
  "FIX002", # Line contains TODO, consider resolving the issue.
  "G004", # Logging statement uses f-string.
  "I001", # Import block is un-sorted or un-formatted.
  "PD008", # Use .loc instead of .at. If speed is important, use NumPy.
  "PD013", # `.melt` is preferred to `.stack`; provides same functionality
  "PGH003", # Use specific rule codes when ignoring type issues.
  "PLR0913", # Too many arguments in function definition.
  "PLR6301", # Method doesn"t use self.
  "PT013", # Found incorrect import of pytest, use simple import pytest instead.
  "Q000", # Single quotes found but double quotes preferred.
  "RUF021", # Parenthesize `a and b` expressions when chaining `and` and `or` together...
  "S101", # Assert used.
  "TC001", # Move application import into a type-checking block.
  "TC002", # Move third-party import into a type-checking block.
  "TC003", # Move standard library import into a type-checking block.
  "TD002", # Missing author in TODO; try: # TODO(<author_name>): ....
  "TD003", # Missing issue link on the line following this TODO.
  "TID252", # Relative imports from parent modules are banned.
]

[tool.ruff.lint.flake8-import-conventions.extend-aliases]
"array_api_extra" = "xpx"
"equinox" = "eqx"
"itertools" = "it"
"jax.numpy" = "jnp"
"jax.random" = "jr"
"jax.scipy.special" = "jss"
"networkx" = "nx"
"numpy.typing" = "npt"
"optimistix" = "optx"
"optype.numpy" = "onp"
"pyqtgraph" = "pg"
"seaborn" = "sns"

[tool.ruff.lint.flake8-errmsg]
max-string-length = 40

[tool.ruff.lint.isort]
combine-as-imports = true

[tool.ruff.lint.pydocstyle]
convention = "google"

[tool.ruff.lint.pylint]
max-args = 8