diff --git a/src/pyrecest/backend_support/_pytorch_trapezoid_numpy_contract.py b/src/pyrecest/backend_support/_pytorch_trapezoid_numpy_contract.py index 24da664d0..ade6c2489 100644 --- a/src/pyrecest/backend_support/_pytorch_trapezoid_numpy_contract.py +++ b/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() diff --git a/tests/backend_support/test_pytorch_triangular_indices_contract.py b/tests/backend_support/test_pytorch_triangular_indices_contract.py new file mode 100644 index 000000000..dc8d5089b --- /dev/null +++ b/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"), + )