array-api-compat/array_api_compat/torch/fft.py at be691e57ae9e596c06af79be6ee63251e91386bc · data-apis/array-api-compat · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from __future__ import annotations

from collections.abc import Sequence
from typing import Literal

import torch  # noqa: F401
import torch.fft

from ._typing import Array
from .._internal import clone_module

__all__ = clone_module("torch.fft", globals())

# Several torch fft functions do not map axes to dim

def fftn(
    x: Array,
    /,
    *,
    s: Sequence[int] | None = None,
    axes: Sequence[int] | None = None,
    norm: Literal["backward", "ortho", "forward"] = "backward",
    **kwargs: object,
) -> Array:
    return torch.fft.fftn(x, s=s, dim=axes, norm=norm, **kwargs)

def ifftn(
    x: Array,
    /,
    *,
    s: Sequence[int] | None = None,
    axes: Sequence[int] | None = None,
    norm: Literal["backward", "ortho", "forward"] = "backward",
    **kwargs: object,
) -> Array:
    return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, **kwargs)

def rfftn(
    x: Array,
    /,
    *,
    s: Sequence[int] | None = None,
    axes: Sequence[int] | None = None,
    norm: Literal["backward", "ortho", "forward"] = "backward",
    **kwargs: object,
) -> Array:
    return torch.fft.rfftn(x, s=s, dim=axes, norm=norm, **kwargs)

def irfftn(
    x: Array,
    /,
    *,
    s: Sequence[int] | None = None,
    axes: Sequence[int] | None = None,
    norm: Literal["backward", "ortho", "forward"] = "backward",
    **kwargs: object,
) -> Array:
    return torch.fft.irfftn(x, s=s, dim=axes, norm=norm, **kwargs)

def fftshift(
    x: Array,
    /,
    *,
    axes: int | Sequence[int] | None = None,
    **kwargs: object,
) -> Array:
    return torch.fft.fftshift(x, dim=axes, **kwargs)

def ifftshift(
    x: Array,
    /,
    *,
    axes: int | Sequence[int] | None = None,
    **kwargs: object,
) -> Array:
    return torch.fft.ifftshift(x, dim=axes, **kwargs)


__all__ += ["fftn", "ifftn", "rfftn", "irfftn", "fftshift", "ifftshift"]

def __dir__() -> list[str]:
    return __all__