pyproject.toml

[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"

[tool.pytest.ini_options]
markers = [
    "multiaccelerator: indicates that a test can make use of and possibly requires multiple accelerators",
    "SlurmMultiNodeGpuTest: mark a test for Slurm multinode GPU nightly CI"
]
filterwarnings = [
    "error",

    # TODO(jakevdp): remove when array_api_tests stabilize
    "default:.*not machine-readable.*:UserWarning",
    "default:Special cases found for .* but none were parsed.*:UserWarning",

    # jax.profiler imports tensorflow.python.profiler.trace internally, which
    # will fail with python 3.12 and some versions of protobuf because the
    # "error" entry above promotes DeprecationWarnings into errors. See also:
    # https://github.com/protocolbuffers/protobuf/issues/12186#issuecomment-1745679358
    "ignore:Type google\\._upb\\._message\\.(Scalar|Message)MapContainer uses PyType_Spec with a metaclass that has custom tp_new\\. This is deprecated and will no longer be allowed in Python 3\\.14\\.:DeprecationWarning",

    # Protobuf version in TF releases may lag behind the ones in code.
    "ignore:Protobuf gencode version \\d+\\.\\d+\\.\\d is exactly one major version older than the runtime version \\d+\\.\\d+\\.\\d at tensorflow/fore/framework/attr_value\\.proto\\. Please update the gencode to avoid compatibility violations in the next runtime release\\.:UserWarning",

    # TODO(b/401588349): Remove this once transparent hugepages are enabled.
    "ignore:Transparent hugepages",

    # TODO(phawkins): Remove this once cloud_tpu_init is removed.
    "ignore:jax.cloud_tpu_init was deprecated",

    # In mpmath 1.4.x, private/deprecated submodules define a module-level
    # __getattr__ that emits a DeprecationWarning on any unknown attribute
    # lookup. When unittest.TestCase.assertWarnsRegex inspects sys.modules
    # and checks getattr(mod, '__warningregistry__', None), it triggers these
    # warnings, which pytest's "error" filter promotes to exceptions.
    "ignore:the (rational|math2).* module is deprecated.*:DeprecationWarning:mpmath.*",

    # NOTE: this is probably not where you want to add code to suppress a
    # warning. Only pytest tests look at this list, whereas Bazel tests also
    # check for warnings and do not check this list. Most likely, you should
    # add a @jtu.ignore_warning decorator to your test instead.
]
doctest_optionflags = [
    "NUMBER",
    "NORMALIZE_WHITESPACE"
]
addopts = "--doctest-glob='*.rst' --ignore='examples/ffi' --import-mode=importlib"

[tool.pyrefly]
python-version = "3.13"
project-includes = [
    "jax/**/*.py*",
]
project-excludes = [
    "jax/_src/internal_test_util/export_back_compat_test_data/*.py",
    "jax/example_libraries/**/*.py",
    "jax/experimental/jax2tf/tests/**/*.py",
]
ignore-missing-imports = [
    "absl.*",
    "cloudpickle.*",
    "filelock.*",
    "flatbuffers.*",
    "flax.*",
    "google.colab.*",
    "hypothesis.*",
    "IPython.*",
    "jax_cuda12_plugin.*",
    "jax_cuda13_plugin.*",
    "jaxlib.*",
    "jraph.*",
    "kubernetes.*",
    "libtpu.*",
    "matplotlib.*",
    "mpi4py.*",
    "nvidia.*",
    "optax.*",
    "opt_einsum.*",
    "portpicker.*",
    "pygments.*",
    "pytest.*",
    "rich.*",
    "setuptools.*",
    "tensorflow.*",
    "tensorflowjs.*",
    "tensorflow_serving.*",
    "web_pdb.*",
    "xprof.*",
    "zstandard.*",
]
untyped-def-behavior = "check-and-infer-return-any"

[tool.pyrefly.errors]
redefinition = "ignore"

[tool.ruff]
preview = true
exclude = [
    ".git",
    "build",
    "__pycache__",
]
line-length = 80
indent-width = 2
target-version = "py312"

[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",
]
select = [
    "B9",
    "C",
    "F",
    "W",
    "YTT",
    "ASYNC",
    "E101",
    "E112",
    "E113",
    "E115",
    "E117",
    "E225",
    "E227",
    "E228",
]

[tool.ruff.lint.mccabe]
max-complexity = 18

[tool.ruff.lint.per-file-ignores]
# F811: Redefinition of unused name.
# F821: Undefined name.
"docs/autodidax.py" = ["F811"]
"docs/pallas/tpu/matmul.ipynb" = ["F811"]
"docs/pallas/tpu/distributed.ipynb" = ["F811"]
"docs/pallas/quickstart.ipynb" = ["F811"]
"docs/notebooks/autodiff_cookbook.ipynb" = ["F811", "F821"]
"docs/notebooks/autodiff_remat.ipynb" = ["F811", "F821"]
"docs/notebooks/Custom_derivative_rules_for_Python_code.ipynb" = ["F811"]
"docs/hijax_custom_derivatives.py" = ["F811"]
"docs/hijax_custom_derivatives.ipynb" = ["F811"]
"docs/jep/9407-type-promotion.ipynb" = ["F811"]
"docs/autodidax.ipynb" = ["F811"]

# This looks like a false positive, but there is no upstream bug.
# TODO(slebedev): Investigate this and remove the ignore.
"*.ipynb" = ["E113"]