[MRG] Load triton before TensorFlow in ot.backend (closes #816)

.github/requirements_doctests.txt

@@ -5,7 +5,7 @@ autograd
 pymanopt
 cvxopt
 scikit-learn
-torch<2.12
+torch
 jax
 jaxlib
 tensorflow; python_version < '3.14'

RELEASES.md

@@ -4,6 +4,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)
 

docs/requirements.txt

@@ -5,7 +5,7 @@ autograd
 pymanopt
 cvxopt
 scikit-learn
-torch<2.12
+torch
 pytest
 torch_geometric
 cvxpy

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

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:]}"
+    )