Fix PyTorch triangular-index NumPy contract

src/pyrecest/backend_support/_pytorch_trapezoid_numpy_contract.py

@@ -113,6 +113,73 @@ def trapezoid(y, x=None, dx=1.0, axis=-1):
         backend.trapezoid = trapezoid
 
 
+def _patch_pytorch_triangular_indices_numpy_contract() -> None:
+    """Patch PyTorch triangular index helpers to follow NumPy's API contract."""
+    try:
+        import pyrecest._backend.pytorch as raw_pytorch  # pylint: disable=import-outside-toplevel
+        import pyrecest.backend as backend  # pylint: disable=import-outside-toplevel
+        import torch  # pylint: disable=import-outside-toplevel
+    except ModuleNotFoundError:  # pragma: no cover - PyTorch backend may be unavailable
+        return
+
+    active_pytorch_backend = getattr(backend, "__backend_name__", None) == "pytorch"
+
+    def _make_triangular_indices(helper_name, torch_index_helper, original_helper):
+        def triangular_indices(
+            n,
+            k=0,
+            m=None,
+            *,
+            dtype=None,
+            device=None,
+            layout=None,
+        ):
+            n = _operator_index(n)
+            k = _operator_index(k)
+            m = n if m is None else _operator_index(m)
+
+            kwargs = {}
+            if dtype is not None:
+                kwargs["dtype"] = dtype
+            if device is not None:
+                kwargs["device"] = device
+            if layout is not None:
+                kwargs["layout"] = layout
+
+            indices = torch_index_helper(
+                row=max(n, 0),
+                col=max(m, 0),
+                offset=k,
+                **kwargs,
+            )
+            return tuple(indices.unbind(0))
+
+        triangular_indices.__name__ = getattr(
+            original_helper, "__name__", helper_name
+        )
+        triangular_indices.__doc__ = getattr(original_helper, "__doc__", None)
+        triangular_indices._pyrecest_numpy_contract = True
+        return triangular_indices
+
+    for helper_name, torch_index_helper in (
+        ("tril_indices", torch.tril_indices),
+        ("triu_indices", torch.triu_indices),
+    ):
+        original_helper = getattr(raw_pytorch, helper_name, None)
+        if original_helper is None:
+            continue
+        if getattr(original_helper, "_pyrecest_numpy_contract", False):
+            if active_pytorch_backend:
+                setattr(backend, helper_name, original_helper)
+            continue
+        helper = _make_triangular_indices(
+            helper_name, torch_index_helper, original_helper
+        )
+        setattr(raw_pytorch, helper_name, helper)
+        if active_pytorch_backend:
+            setattr(backend, helper_name, helper)
+
+
 def _patch_rectangular_pytorch_triangular_vector_contract() -> None:
     """Patch PyTorch triangular vector helpers for rectangular matrices."""
     try:
@@ -204,5 +271,6 @@ def triangular_to_vec(x, k=0):
             setattr(backend, helper_name, helper)
 
 
+_patch_pytorch_triangular_indices_numpy_contract()
 _patch_rectangular_pytorch_triangular_vector_contract()
 _patch_rectangular_jax_triangular_vector_contract()

tests/backend_support/test_pytorch_triangular_indices_contract.py

@@ -0,0 +1,86 @@
+import importlib.util
+import os
+import subprocess
+import sys
+
+import pytest
+
+
+def _backend_test_env(backend_name):
+    env = os.environ.copy()
+    env["PYRECEST_BACKEND"] = backend_name
+    src_path = os.path.abspath("src")
+    env["PYTHONPATH"] = (
+        src_path
+        if not env.get("PYTHONPATH")
+        else os.pathsep.join([src_path, env["PYTHONPATH"]])
+    )
+    return env
+
+
+@pytest.mark.backend_portable
+def test_pytorch_triangular_indices_follow_numpy_contract():
+    if importlib.util.find_spec("torch") is None:
+        pytest.skip("torch is not installed")
+
+    code = """
+import numpy as np
+import pyrecest.backend as backend
+import pyrecest._backend.pytorch as pytorch_backend
+
+for module in (backend, pytorch_backend):
+    for helper_name in ("tril_indices", "triu_indices"):
+        helper = getattr(module, helper_name)
+        numpy_helper = getattr(np, helper_name)
+
+        for args, kwargs in (
+            ((3,), {}),
+            ((3,), {"k": 1}),
+            ((2,), {"k": -1, "m": 4}),
+            ((-1,), {}),
+            ((3,), {"m": -1}),
+        ):
+            result = helper(*args, **kwargs)
+            expected = numpy_helper(*args, **kwargs)
+            assert isinstance(result, tuple)
+            assert len(result) == 2
+            for actual_axis, expected_axis in zip(result, expected):
+                assert module.to_numpy(actual_axis).tolist() == expected_axis.tolist()
+"""
+    subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        env=_backend_test_env("pytorch"),
+    )
+
+
+@pytest.mark.backend_portable
+def test_raw_pytorch_triangular_indices_are_patched_with_numpy_public_backend():
+    if importlib.util.find_spec("torch") is None:
+        pytest.skip("torch is not installed")
+
+    code = """
+import numpy as np
+import pyrecest.backend as backend
+import pyrecest._backend.pytorch as pytorch_backend
+import torch
+
+assert getattr(backend, "__backend_name__", None) == "numpy"
+
+rows, cols = pytorch_backend.tril_indices(3)
+expected_rows, expected_cols = np.tril_indices(3)
+assert pytorch_backend.to_numpy(rows).tolist() == expected_rows.tolist()
+assert pytorch_backend.to_numpy(cols).tolist() == expected_cols.tolist()
+
+rows, cols = pytorch_backend.triu_indices(2, k=1, m=4, dtype=torch.int32)
+expected_rows, expected_cols = np.triu_indices(2, k=1, m=4)
+assert rows.dtype == torch.int32
+assert cols.dtype == torch.int32
+assert pytorch_backend.to_numpy(rows).tolist() == expected_rows.tolist()
+assert pytorch_backend.to_numpy(cols).tolist() == expected_cols.tolist()
+"""
+    subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        env=_backend_test_env("numpy"),
+    )