Repository: scipy/scipy
File: scipy/fft/tests/test_basic.py (Python)

Code before:
            expect * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
        )
        xp_assert_close(fft.ihfftn(x, norm="forward"), expect * (30 * 20 * 10))

    def _check_axes(self, op, xp):
        dtype = get_expected_input_dtype(op, xp)
        x = xp.asarray(random((30, 20, 10)), dtype=dtype)
        axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]

        for a in axes:
            op_tr = op(xp.permute_dims(x, axes=a))
            tr_op = xp.permute_dims(op(x, axes=a), axes=a)
            xp_assert_close(op_tr, tr_op)

    @make_xp_test_case(fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn)
    @pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn])
    def test_axes_standard(self, op, xp):
        self._check_axes(op, xp)

    @make_xp_test_case(fft.hfftn, fft.ihfftn)
    @pytest.mark.parametrize("op", [fft.hfftn, fft.ihfftn])
    def test_axes_non_standard(self, op, xp):
        self._check_axes(op, xp)

    @make_xp_test_case(fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn)
    @pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn])
    def test_axes_subset_with_shape_standard(self, op, xp):
        dtype = get_expected_input_dtype(op, xp)
        x = xp.asarray(random((16, 8, 4)), dtype=dtype)
        axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]

        for a in axes:
            # different shape on the first two axes
            shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
                           for ax in range(x.ndim)])
            # transform only the first two axes
            op_tr = op(xp.permute_dims(x, axes=a),
                       s=shape[:2], axes=(0, 1))
            tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]),
                                         axes=a)
            xp_assert_close(op_tr, tr_op)

    @make_xp_test_case(fft.fft2, fft.ifft2, fft.rfft2, fft.irfft2, fft.hfft2,
                       fft.ihfft2, fft.hfftn, fft.ihfftn)
    @pytest.mark.parametrize("op", [fft.fft2, fft.ifft2, fft.rfft2, fft.irfft2,
                                    fft.hfft2, fft.ihfft2, fft.hfftn, fft.ihfftn])
    def test_axes_subset_with_shape_non_standard(self, op, xp):
        dtype = get_expected_input_dtype(op, xp)
        x = xp.asarray(random((16, 8, 4)), dtype=dtype)
        axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]


Review comment (suggestion):
One thing I can't help noticing is that @steppi used `make_xp_test_param` to parametrize over array api capable functions: https://github.com/scipy/scipy/pull/24024#discussion_r2546403287

It'd be best to converge on a uniform way to handle these, whichever it is (and I don't have a foggiest as to which one).

Code after:
            expect * xp.sqrt(xp.asarray(30 * 20 * 10, dtype=xp.float64))
        )
        xp_assert_close(fft.ihfftn(x, norm="forward"), expect * (30 * 20 * 10))

    def _check_axes(self, op, xp):
        dtype = get_expected_input_dtype(op, xp)
        x = xp.asarray(random((30, 20, 10)), dtype=dtype)
        axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]

        for a in axes:
            op_tr = op(xp.permute_dims(x, axes=a))
            tr_op = xp.permute_dims(op(x, axes=a), axes=a)
            xp_assert_close(op_tr, tr_op)

    @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fftn),
                                    make_xp_pytest_param(fft.ifftn),
                                    make_xp_pytest_param(fft.rfftn),
                                    make_xp_pytest_param(fft.irfftn)])
    def test_axes_standard(self, op, xp):
        self._check_axes(op, xp)

    @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.hfftn),
                                    make_xp_pytest_param(fft.ihfftn)])
    def test_axes_non_standard(self, op, xp):
        self._check_axes(op, xp)

    @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fftn),
                                    make_xp_pytest_param(fft.ifftn),
                                    make_xp_pytest_param(fft.rfftn),
                                    make_xp_pytest_param(fft.irfftn)])
    def test_axes_subset_with_shape_standard(self, op, xp):
        dtype = get_expected_input_dtype(op, xp)
        x = xp.asarray(random((16, 8, 4)), dtype=dtype)
        axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]

        for a in axes:
            # different shape on the first two axes
            shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
                           for ax in range(x.ndim)])
            # transform only the first two axes
            op_tr = op(xp.permute_dims(x, axes=a),
                       s=shape[:2], axes=(0, 1))
            tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]),
                                         axes=a)
            xp_assert_close(op_tr, tr_op)

    @pytest.mark.parametrize("op", [make_xp_pytest_param(fft.fft2),
                                    make_xp_pytest_param(fft.ifft2),
                                    make_xp_pytest_param(fft.rfft2),
                                    make_xp_pytest_param(fft.irfft2),
                                    make_xp_pytest_param(fft.hfft2),