diff --git a/src/pyrecest/_backend/pytorch/fft.py b/src/pyrecest/_backend/pytorch/fft.py index 539ca4e52..e0d037a38 100644 --- a/src/pyrecest/_backend/pytorch/fft.py +++ b/src/pyrecest/_backend/pytorch/fft.py @@ -124,6 +124,26 @@ def _normalize_fft_shape_args(args, kwargs): return args, kwargs +def _normalize_fft_dim_args(args, kwargs, position, normalizer): + """Normalize a positional or keyword ``dim`` FFT argument.""" + if position is not None and len(args) > position: + args = list(args) + args[position] = normalizer(args[position]) + return tuple(args), kwargs + if "dim" not in kwargs: + return args, kwargs + kwargs = dict(kwargs) + kwargs["dim"] = normalizer(kwargs["dim"]) + return args, kwargs + + +def _get_fft_dim_arg(args, kwargs, position): + """Return the effective positional or keyword ``dim`` FFT argument.""" + if position is not None and len(args) > position: + return args[position] + return kwargs.get("dim") + + def _with_dim_alias(kwargs, alias, func_name, *, none_alias_means_default=True): if alias not in kwargs: return kwargs @@ -153,6 +173,7 @@ def _wrap_arraylike_fft( *, func_name, dim_alias=None, + dim_arg_position=None, empty_dim_is_noop=False, normalize_scalar_dim=False, normalize_dim_sequence=False, @@ -171,18 +192,21 @@ def fft_func(value, *args, **kwargs): ) if validate_real_length: _validate_real_fft_length_args(args, kwargs) - if normalize_scalar_dim and "dim" in kwargs: - kwargs = dict(kwargs) - kwargs["dim"] = _normalize_single_fft_dim(kwargs["dim"]) - if normalize_dim_sequence and "dim" in kwargs: - kwargs = dict(kwargs) - kwargs["dim"] = _normalize_fft_dim_sequence(kwargs["dim"]) + if normalize_scalar_dim: + args, kwargs = _normalize_fft_dim_args( + args, kwargs, dim_arg_position, _normalize_single_fft_dim + ) + if normalize_dim_sequence: + args, kwargs = _normalize_fft_dim_args( + args, kwargs, dim_arg_position, _normalize_fft_dim_sequence + ) if normalize_shape_sequence: args, kwargs = _normalize_fft_shape_args(args, kwargs) value = _as_fft_tensor(value) + dim = _get_fft_dim_arg(args, kwargs, dim_arg_position) if ( empty_dim_is_noop - and _is_empty_dim(kwargs.get("dim")) + and _is_empty_dim(dim) and _empty_dim_noop_is_valid(args, kwargs) ): return value @@ -195,6 +219,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.rfft, func_name="rfft", dim_alias="axis", + dim_arg_position=1, normalize_scalar_dim=True, validate_real_length=True, none_alias_means_default=False, @@ -203,6 +228,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.irfft, func_name="irfft", dim_alias="axis", + dim_arg_position=1, normalize_scalar_dim=True, validate_real_length=True, none_alias_means_default=False, @@ -211,6 +237,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.fftshift, func_name="fftshift", dim_alias="axes", + dim_arg_position=0, empty_dim_is_noop=True, normalize_dim_sequence=True, ) @@ -218,6 +245,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.ifftshift, func_name="ifftshift", dim_alias="axes", + dim_arg_position=0, empty_dim_is_noop=True, normalize_dim_sequence=True, ) @@ -225,6 +253,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.fftn, func_name="fftn", dim_alias="axes", + dim_arg_position=1, empty_dim_is_noop=True, normalize_dim_sequence=True, normalize_shape_sequence=True, @@ -233,6 +262,7 @@ def fft_func(value, *args, **kwargs): _torch.fft.ifftn, func_name="ifftn", dim_alias="axes", + dim_arg_position=1, empty_dim_is_noop=True, normalize_dim_sequence=True, normalize_shape_sequence=True, diff --git a/tests/backend_support/test_pytorch_fft_positional_axes_contract.py b/tests/backend_support/test_pytorch_fft_positional_axes_contract.py new file mode 100644 index 000000000..5ad9c0820 --- /dev/null +++ b/tests/backend_support/test_pytorch_fft_positional_axes_contract.py @@ -0,0 +1,83 @@ +import importlib.util +import os +import subprocess +import sys + +import pytest + + +@pytest.mark.backend_portable +def test_pytorch_fft_positional_axes_accept_numpy_integer_arrays(): + if importlib.util.find_spec("torch") is None: + pytest.skip("torch is not installed") + + env = os.environ.copy() + env["PYRECEST_BACKEND"] = "pytorch" + src_path = os.path.abspath("src") + env["PYTHONPATH"] = ( + src_path + if not env.get("PYTHONPATH") + else os.pathsep.join([src_path, env["PYTHONPATH"]]) + ) + + code = """ +import numpy as np +import numpy.testing as npt +import pyrecest.backend as backend + +values_np = np.array([[0.0, 1.0], [2.0, 3.0]]) +values = backend.array(values_np) + +for axis in (np.array(0), np.int64(0)): + npt.assert_allclose( + backend.to_numpy(backend.fft.rfft(values, None, axis)), + np.fft.rfft(values_np, None, axis), + ) + npt.assert_allclose( + backend.to_numpy(backend.fft.irfft(values, None, axis)), + np.fft.irfft(values_np, None, axis), + ) + +axes_cases = ( + np.array([0]), + np.array([0, 1]), + [np.array(0)], + (np.array(1),), +) +for axes in axes_cases: + npt.assert_array_equal( + backend.to_numpy(backend.fft.fftshift(values, axes)), + np.fft.fftshift(values_np, axes), + ) + npt.assert_array_equal( + backend.to_numpy(backend.fft.ifftshift(values, axes)), + np.fft.ifftshift(values_np, axes), + ) + npt.assert_allclose( + backend.to_numpy(backend.fft.fftn(values, None, axes)), + np.fft.fftn(values_np, None, axes), + ) + npt.assert_allclose( + backend.to_numpy(backend.fft.ifftn(values, None, axes)), + np.fft.ifftn(values_np, None, axes), + ) + +for axes in ((), []): + npt.assert_array_equal( + backend.to_numpy(backend.fft.fftshift(values, axes)), + values_np, + ) + npt.assert_array_equal( + backend.to_numpy(backend.fft.ifftshift(values, axes)), + values_np, + ) + npt.assert_array_equal( + backend.to_numpy(backend.fft.fftn(values, None, axes)), + values_np, + ) + npt.assert_array_equal( + backend.to_numpy(backend.fft.ifftn(values, None, axes)), + values_np, + ) +""" + subprocess.run([sys.executable, "-c", code], check=True, env=env)