diff --git a/.github/requirements_doctests.txt b/.github/requirements_doctests.txt index cae6a079d..7c1aa4990 100644 --- a/.github/requirements_doctests.txt +++ b/.github/requirements_doctests.txt @@ -5,7 +5,7 @@ autograd pymanopt cvxopt scikit-learn -torch<2.12 +torch jax jaxlib tensorflow; python_version < '3.14' diff --git a/RELEASES.md b/RELEASES.md index 040b35af3..33be0db8e 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,7 @@ #### Closed issues +- Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816) - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840) - `ot.dr.fda` and `ot.dr.wda` no longer modify the input array `X` in place (PR #840) diff --git a/docs/requirements.txt b/docs/requirements.txt index beb66b6bc..021e833ca 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -5,7 +5,7 @@ autograd pymanopt cvxopt scikit-learn -torch<2.12 +torch pytest torch_geometric cvxpy diff --git a/ot/backend.py b/ot/backend.py index 7749c948d..af622734d 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -86,6 +86,7 @@ # # License: MIT License +import importlib.util import os import time import warnings @@ -108,6 +109,21 @@ import torch torch_type = torch.Tensor + + # Load triton before TensorFlow is imported below. Both triton and + # TensorFlow ship a statically linked LLVM, and loading triton's + # libtriton.so into a process that has already imported TensorFlow + # segfaults inside dlopen. torch only imports triton lazily, on the + # first use of a feature that needs it (constructing an optimizer is + # enough), which would otherwise happen after TensorFlow is loaded. + # See https://github.com/PythonOT/POT/issues/816 + if not os.environ.get(DISABLE_TF_KEY, False) and ( + importlib.util.find_spec("tensorflow") is not None + ): + try: + import triton # noqa: F401 + except ImportError: + pass except ImportError: torch = False torch_type = float diff --git a/test/test_backend.py b/test/test_backend.py index fe6af9c67..4df918140 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -6,6 +6,10 @@ # # License: MIT License +import importlib.util +import subprocess +import sys + import numpy as np import pytest from numpy.testing import assert_array_almost_equal_nulp @@ -914,3 +918,27 @@ def test_get_backend_none(): assert str(nx) == "numpy" with pytest.raises(ValueError): get_backend(None, None) + + +@pytest.mark.skipif( + not torch or not tf or importlib.util.find_spec("triton") is None, + reason="Requires torch, tensorflow and triton installed together", +) +def test_torch_optimizer_after_tensorflow_import(): + """Non-regression test for issue #816. + + Building a torch optimizer makes torch import triton lazily. If TensorFlow + was imported first, loading libtriton.so segfaults the interpreter, so this + has to run in a subprocess. + """ + code = ( + "import ot.backend\n" + "import torch\n" + "x = torch.zeros(3, requires_grad=True)\n" + "torch.optim.SGD([x], lr=0.1)\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True) + assert result.returncode == 0, ( + f"interpreter died with returncode {result.returncode}: " + f"{result.stderr.decode(errors='replace')[-2000:]}" + )