philtorch 0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. philtorch-0.1/LICENSE +21 -0
  2. philtorch-0.1/MANIFEST.in +2 -0
  3. philtorch-0.1/PKG-INFO +163 -0
  4. philtorch-0.1/README.md +138 -0
  5. philtorch-0.1/philtorch/VERSION.txt +1 -0
  6. philtorch-0.1/philtorch/__init__.py +13 -0
  7. philtorch-0.1/philtorch/_version.py +18 -0
  8. philtorch-0.1/philtorch/core.py +0 -0
  9. philtorch-0.1/philtorch/csrc/host_lti_recur.cpp +85 -0
  10. philtorch-0.1/philtorch/csrc/host_lti_recurN.cpp +114 -0
  11. philtorch-0.1/philtorch/csrc/host_recur2.cpp +171 -0
  12. philtorch-0.1/philtorch/csrc/host_recurN.cpp +287 -0
  13. philtorch-0.1/philtorch/csrc/recur2.cuh +77 -0
  14. philtorch-0.1/philtorch/lpv/__init__.py +12 -0
  15. philtorch-0.1/philtorch/lpv/filtering.py +370 -0
  16. philtorch-0.1/philtorch/lpv/recur.py +101 -0
  17. philtorch-0.1/philtorch/lpv/ssm.py +548 -0
  18. philtorch-0.1/philtorch/lpv/utils.py +33 -0
  19. philtorch-0.1/philtorch/lti/__init__.py +16 -0
  20. philtorch-0.1/philtorch/lti/filtering.py +484 -0
  21. philtorch-0.1/philtorch/lti/recur.py +192 -0
  22. philtorch-0.1/philtorch/lti/ssm.py +798 -0
  23. philtorch-0.1/philtorch/mat.py +166 -0
  24. philtorch-0.1/philtorch/poly.py +36 -0
  25. philtorch-0.1/philtorch/prototype/poly.py +181 -0
  26. philtorch-0.1/philtorch/prototype/signal.py +373 -0
  27. philtorch-0.1/philtorch/utils.py +27 -0
  28. philtorch-0.1/philtorch.egg-info/PKG-INFO +163 -0
  29. philtorch-0.1/philtorch.egg-info/SOURCES.txt +43 -0
  30. philtorch-0.1/philtorch.egg-info/dependency_links.txt +1 -0
  31. philtorch-0.1/philtorch.egg-info/not-zip-safe +1 -0
  32. philtorch-0.1/philtorch.egg-info/requires.txt +8 -0
  33. philtorch-0.1/philtorch.egg-info/top_level.txt +2 -0
  34. philtorch-0.1/pyproject.toml +71 -0
  35. philtorch-0.1/setup.cfg +4 -0
  36. philtorch-0.1/setup.py +60 -0
  37. philtorch-0.1/tests/test_grad.py +230 -0
  38. philtorch-0.1/tests/test_lpv_filters.py +419 -0
  39. philtorch-0.1/tests/test_lpv_lfilter.py +99 -0
  40. philtorch-0.1/tests/test_lpv_ssm.py +281 -0
  41. philtorch-0.1/tests/test_lti_comb.py +44 -0
  42. philtorch-0.1/tests/test_lti_lfilter.py +425 -0
  43. philtorch-0.1/tests/test_lti_ssm.py +196 -0
  44. philtorch-0.1/tests/test_recur_ext.py +95 -0
  45. philtorch-0.1/tests/test_scan.py +52 -0
philtorch-0.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Chin-Yun Yu
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,2 @@
1
+ recursive-include philtorch *.txt
2
+ recursive-include philtorch *.cuh
philtorch-0.1/PKG-INFO ADDED
@@ -0,0 +1,163 @@
1
+ Metadata-Version: 2.4
2
+ Name: philtorch
3
+ Version: 0.1
4
+ Summary: A PyTorch library for time-varying IIR filters.
5
+ Author-email: Chin-Yun Yu <chin-yun.yu@qmul.ac.uk>
6
+ Maintainer-email: Chin-Yun Yu <chin-yun.yu@qmul.ac.uk>
7
+ License-Expression: MIT
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.10
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Requires-Python: >=3.10
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: torch>=2.0.0
18
+ Requires-Dist: torchlpc>=0.7.2
19
+ Provides-Extra: dev
20
+ Requires-Dist: pytest>=6.0; extra == "dev"
21
+ Requires-Dist: scipy; extra == "dev"
22
+ Requires-Dist: numpy; extra == "dev"
23
+ Requires-Dist: pytest-cov; extra == "dev"
24
+ Dynamic: license-file
25
+
26
+ # PhilTorch
27
+
28
+ A PyTorch package for fast automatic differentiation of discrete time linear filters.
29
+
30
+ Our principle design goals are:
31
+
32
+ - Provide fast and differentiable version of `scipy.signal.*` functions.
33
+ - Focus on time-domain implementation without using FFT.
34
+ - Support batch processing, parameter-varying filters, and GPU acceleration.
35
+ - Pure functional implementation and no stateful objects.
36
+
37
+ ## Installation
38
+
39
+ ### Stable release
40
+ ```bash
41
+ pip install philtorch
42
+ ```
43
+
44
+ ### Development version
45
+ ```bash
46
+ pip install -i https://test.pypi.org/simple/ philtorch
47
+ # or
48
+ pip install git+https://github.com/yoyolicoris/philtorch.git
49
+ ```
50
+ > **_Note:_**
51
+ > - The installation process compiles C++/CUDA extensions, so make sure you have a working C++ compiler and CUDA toolkit (if you want to use GPU acceleration) installed.
52
+ > - We recommend using `--no-build-isolation` flag to avoid potential issues with building the package in an isolated environment, especially when installing with CUDA support.
53
+
54
+ ## Module overview
55
+
56
+ - `philtorch`: Root module.
57
+ - `lpv`: Functions under it are for linear parameter-varying filters.
58
+ - `fir`:
59
+ - Finite Impulse Response filters.
60
+ - `allpole`:
61
+ - All-pole filters.
62
+ - `lfilter`:
63
+ - Parameter-varying version of `scipy.signal.lfilter`. It supports not only transposed direct form II but also transposed direct form I, direct form I, and direct form II structures.
64
+ - `state_space`:
65
+ - Parameter-varying state-space models.
66
+ - `state_space_recursion`:
67
+ - The core recursion function for state-space models.
68
+ - `linear_recurrence`:
69
+ - A linear recurrence function with scalar coefficients.
70
+ - `lti`: Functions under it are for linear time-invariant filters.
71
+ - `fir`:
72
+ - Finite Impulse Response filters.
73
+ - `lfilter`:
74
+ - A differentiable version of `scipy.signal.lfilter`. It supports not only transposed direct form II but also transposed direct form I, direct form I, and direct form II structures.
75
+ - `filtfilt`:
76
+ - A differentiable version of `scipy.signal.filtfilt`.
77
+ - `lfilter_zi`:
78
+ - A differentiable version of `scipy.signal.lfilter_zi`.
79
+ - `lfiltic`:
80
+ - A differentiable version of `scipy.signal.lfiltic`.
81
+ - `state_space`:
82
+ - State-space models.
83
+ - `diag_state_space`:
84
+ - State-space models with diagonalisable state matrix.
85
+ - `state_space_recursion`:
86
+ - The core recursion function for state-space models.
87
+ - `linear_recurrence`:
88
+ - A linear recurrence function with scalar coefficients.
89
+ - `utils`: Utility functions.
90
+ - `mat`: Matrix operations.
91
+ - `poly`: Polynomial operations.
92
+
93
+ For detailed API reference, please refer to the docstring of each function.
94
+
95
+
96
+
97
+ ## Examples
98
+
99
+ ### Recreating `scipy.signal.lfilter` example
100
+
101
+ ```python
102
+ import torch
103
+ from philtorch.lti import lfilter, lfilter_zi, filtfilt
104
+ from scipy.signal import butter
105
+
106
+ x = torch.randn(201)
107
+
108
+ b_np, a_np = butter(3, 0.05)
109
+ # note that in philtorch a_0 is always 1
110
+ b_np /= a_np[0]
111
+ a_np = a_np[1:] / a_np[0]
112
+ b, a = torch.from_numpy(b_np), torch.from_numpy(a_np)
113
+
114
+ # note that the position of a and b are swapped compared to scipy
115
+ zi = lfilter_zi(a, b)
116
+
117
+ z, _ = lfilter(b, a, x, zi=zi * x[0])
118
+ z2 = filtfilt(b, a, x)
119
+ ```
120
+
121
+ If `lfilter` is imported from `philtorch.lpv`, it can also handle parameter-varying filters, where `a` and `b` are at least 2D tensors with an additional time dimension.
122
+
123
+ ### Computing the first 10 Fibonacci numbers using `state_space`
124
+
125
+ The function `philtorch.lti.state_space` compute the following recursion:
126
+
127
+ ```math
128
+ \begin{aligned}
129
+ \mathbf{h}_{n+1} &= \mathbf{A} \mathbf{h}_n + \mathbf{B} \mathbf{x}_n \\
130
+ \mathbf{y}_n &= \mathbf{C} \mathbf{h}_n + \mathbf{D} \mathbf{x}_n
131
+ \end{aligned}
132
+ ```
133
+
134
+ We can use it to compute the Fibonacci numbers by setting:
135
+
136
+ ```math
137
+ \begin{aligned}
138
+ \mathbf{A} = \begin{bmatrix} 1 & 1 \\ 1 & 0 \end{bmatrix}, \quad
139
+ \mathbf{C} = \begin{bmatrix} 1 & 0 \end{bmatrix}, \quad
140
+ \mathbf{B} = \mathbf{D} = 0, \\
141
+ \mathbf{h}_0 = \begin{bmatrix} 1 \\ 0 \end{bmatrix}
142
+ \end{aligned}
143
+ ```
144
+
145
+ ```python
146
+ import torch
147
+ from philtorch.lti import state_space
148
+
149
+ A = torch.tensor([[1, 1], [1, 0]])
150
+ C = torch.tensor([1, 0])
151
+ x = torch.zeros(1, 10).long()
152
+ h0 = torch.tensor([1, 0])
153
+ y, _ = state_space(A, x, C=C, zi=h0)
154
+ print(y)
155
+ ```
156
+ ```
157
+ tensor([[ 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]])
158
+ ```
159
+ The result is the first 10 Fibonacci numbers, which has the following recursion relation:
160
+
161
+ ```math
162
+ F_n = F_{n-1} + F_{n-2}, \quad F_0 = 1, \quad F_1 = 1
163
+ ```
@@ -0,0 +1,138 @@
1
+ # PhilTorch
2
+
3
+ A PyTorch package for fast automatic differentiation of discrete time linear filters.
4
+
5
+ Our principle design goals are:
6
+
7
+ - Provide fast and differentiable version of `scipy.signal.*` functions.
8
+ - Focus on time-domain implementation without using FFT.
9
+ - Support batch processing, parameter-varying filters, and GPU acceleration.
10
+ - Pure functional implementation and no stateful objects.
11
+
12
+ ## Installation
13
+
14
+ ### Stable release
15
+ ```bash
16
+ pip install philtorch
17
+ ```
18
+
19
+ ### Development version
20
+ ```bash
21
+ pip install -i https://test.pypi.org/simple/ philtorch
22
+ # or
23
+ pip install git+https://github.com/yoyolicoris/philtorch.git
24
+ ```
25
+ > **_Note:_**
26
+ > - The installation process compiles C++/CUDA extensions, so make sure you have a working C++ compiler and CUDA toolkit (if you want to use GPU acceleration) installed.
27
+ > - We recommend using `--no-build-isolation` flag to avoid potential issues with building the package in an isolated environment, especially when installing with CUDA support.
28
+
29
+ ## Module overview
30
+
31
+ - `philtorch`: Root module.
32
+ - `lpv`: Functions under it are for linear parameter-varying filters.
33
+ - `fir`:
34
+ - Finite Impulse Response filters.
35
+ - `allpole`:
36
+ - All-pole filters.
37
+ - `lfilter`:
38
+ - Parameter-varying version of `scipy.signal.lfilter`. It supports not only transposed direct form II but also transposed direct form I, direct form I, and direct form II structures.
39
+ - `state_space`:
40
+ - Parameter-varying state-space models.
41
+ - `state_space_recursion`:
42
+ - The core recursion function for state-space models.
43
+ - `linear_recurrence`:
44
+ - A linear recurrence function with scalar coefficients.
45
+ - `lti`: Functions under it are for linear time-invariant filters.
46
+ - `fir`:
47
+ - Finite Impulse Response filters.
48
+ - `lfilter`:
49
+ - A differentiable version of `scipy.signal.lfilter`. It supports not only transposed direct form II but also transposed direct form I, direct form I, and direct form II structures.
50
+ - `filtfilt`:
51
+ - A differentiable version of `scipy.signal.filtfilt`.
52
+ - `lfilter_zi`:
53
+ - A differentiable version of `scipy.signal.lfilter_zi`.
54
+ - `lfiltic`:
55
+ - A differentiable version of `scipy.signal.lfiltic`.
56
+ - `state_space`:
57
+ - State-space models.
58
+ - `diag_state_space`:
59
+ - State-space models with diagonalisable state matrix.
60
+ - `state_space_recursion`:
61
+ - The core recursion function for state-space models.
62
+ - `linear_recurrence`:
63
+ - A linear recurrence function with scalar coefficients.
64
+ - `utils`: Utility functions.
65
+ - `mat`: Matrix operations.
66
+ - `poly`: Polynomial operations.
67
+
68
+ For detailed API reference, please refer to the docstring of each function.
69
+
70
+
71
+
72
+ ## Examples
73
+
74
+ ### Recreating `scipy.signal.lfilter` example
75
+
76
+ ```python
77
+ import torch
78
+ from philtorch.lti import lfilter, lfilter_zi, filtfilt
79
+ from scipy.signal import butter
80
+
81
+ x = torch.randn(201)
82
+
83
+ b_np, a_np = butter(3, 0.05)
84
+ # note that in philtorch a_0 is always 1
85
+ b_np /= a_np[0]
86
+ a_np = a_np[1:] / a_np[0]
87
+ b, a = torch.from_numpy(b_np), torch.from_numpy(a_np)
88
+
89
+ # note that the position of a and b are swapped compared to scipy
90
+ zi = lfilter_zi(a, b)
91
+
92
+ z, _ = lfilter(b, a, x, zi=zi * x[0])
93
+ z2 = filtfilt(b, a, x)
94
+ ```
95
+
96
+ If `lfilter` is imported from `philtorch.lpv`, it can also handle parameter-varying filters, where `a` and `b` are at least 2D tensors with an additional time dimension.
97
+
98
+ ### Computing the first 10 Fibonacci numbers using `state_space`
99
+
100
+ The function `philtorch.lti.state_space` compute the following recursion:
101
+
102
+ ```math
103
+ \begin{aligned}
104
+ \mathbf{h}_{n+1} &= \mathbf{A} \mathbf{h}_n + \mathbf{B} \mathbf{x}_n \\
105
+ \mathbf{y}_n &= \mathbf{C} \mathbf{h}_n + \mathbf{D} \mathbf{x}_n
106
+ \end{aligned}
107
+ ```
108
+
109
+ We can use it to compute the Fibonacci numbers by setting:
110
+
111
+ ```math
112
+ \begin{aligned}
113
+ \mathbf{A} = \begin{bmatrix} 1 & 1 \\ 1 & 0 \end{bmatrix}, \quad
114
+ \mathbf{C} = \begin{bmatrix} 1 & 0 \end{bmatrix}, \quad
115
+ \mathbf{B} = \mathbf{D} = 0, \\
116
+ \mathbf{h}_0 = \begin{bmatrix} 1 \\ 0 \end{bmatrix}
117
+ \end{aligned}
118
+ ```
119
+
120
+ ```python
121
+ import torch
122
+ from philtorch.lti import state_space
123
+
124
+ A = torch.tensor([[1, 1], [1, 0]])
125
+ C = torch.tensor([1, 0])
126
+ x = torch.zeros(1, 10).long()
127
+ h0 = torch.tensor([1, 0])
128
+ y, _ = state_space(A, x, C=C, zi=h0)
129
+ print(y)
130
+ ```
131
+ ```
132
+ tensor([[ 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]])
133
+ ```
134
+ The result is the first 10 Fibonacci numbers, which has the following recursion relation:
135
+
136
+ ```math
137
+ F_n = F_{n-1} + F_{n-2}, \quad F_0 = 1, \quad F_1 = 1
138
+ ```
@@ -0,0 +1 @@
1
+ 0.1
@@ -0,0 +1,13 @@
1
+ from pathlib import Path
2
+ import warnings
3
+ import torch
4
+
5
+ try:
6
+ from . import _C
7
+
8
+ EXTENSION_LOADED = True
9
+ except ImportError:
10
+ EXTENSION_LOADED = False
11
+ warnings.warn("Custom extension not loaded.")
12
+
13
+ __version__ = Path(__file__).parent.joinpath("VERSION.txt").read_text()
@@ -0,0 +1,18 @@
1
+ import re
2
+
3
+
4
+ def format_branch_name(name):
5
+ # "(fix|feat)/issue-name" or CICD's branch "HEAD"
6
+ pattern = re.compile("^((fix|feat)\/(?P<branch>.+))|((head|HEAD))")
7
+
8
+ match = pattern.search(name)
9
+ if match:
10
+ return f"dev+{match.group(0)}" # => dev+"(fix|feat)/issue-name"
11
+
12
+ # function is called even if branch name is not used in a current template
13
+ # just left properly named branches intact
14
+ if name in ["master", "dev", "main"]:
15
+ return name
16
+
17
+ # fail in case of wrong branch names like "bugfix/issue-unknown"
18
+ raise ValueError(f"Wrong branch name: {name}")
File without changes
@@ -0,0 +1,85 @@
1
+ #include <torch/script.h>
2
+ #include <torch/torch.h>
3
+
4
+ template <typename scalar_t>
5
+ void host_lti_batch_linear_recurrence(int B, int T,
6
+ const scalar_t *a,
7
+ scalar_t *out)
8
+ {
9
+ // Process each batch in parallel
10
+ #pragma omp parallel for
11
+ for (int b = 0; b < B; ++b)
12
+ {
13
+ const scalar_t a_b = a[b];
14
+ scalar_t *out_b = out + b * T;
15
+
16
+ // t loop
17
+ for (int t = 1; t < T; ++t)
18
+ {
19
+ out_b[t] += a_b * out_b[t - 1];
20
+ }
21
+ }
22
+ }
23
+
24
+ template <typename scalar_t>
25
+ void host_lti_shared_linear_recurrence(int B, int T,
26
+ const scalar_t a,
27
+ scalar_t *out)
28
+ {
29
+ // Process each batch in parallel
30
+ #pragma omp parallel for
31
+ for (int b = 0; b < B; ++b)
32
+ {
33
+ scalar_t *out_b = out + b * T;
34
+ // t loop
35
+ for (int t = 1; t < T; ++t)
36
+ {
37
+ out_b[t] += a * out_b[t - 1];
38
+ }
39
+ }
40
+ }
41
+
42
+ at::Tensor lti_recur_cpu_impl(const at::Tensor &a,
43
+ const at::Tensor &zi, const at::Tensor &x)
44
+ {
45
+ TORCH_CHECK(zi.scalar_type() == x.scalar_type(),
46
+ "zi must have the same scalar type as input");
47
+ TORCH_CHECK(a.scalar_type() == x.scalar_type(),
48
+ "A must have the same scalar type as input");
49
+ TORCH_CHECK(a.dim() <= 1, "A must be a vector or a scalar");
50
+ TORCH_CHECK(zi.dim() == 1, "zi must be a vector");
51
+
52
+ auto n_steps = x.size(1) + 1; // +1 for the initial state
53
+ auto n_batches = x.size(0);
54
+ auto output =
55
+ at::cat({zi.unsqueeze(1), x}, 1).contiguous();
56
+ auto a_contiguous = a.contiguous();
57
+
58
+ if (a.dim() == 1 && a.numel() == n_batches)
59
+ {
60
+
61
+ AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
62
+ x.scalar_type(), "host_lti_batch_linear_recurrence", [&]
63
+ { host_lti_batch_linear_recurrence<scalar_t>(
64
+ n_batches, n_steps,
65
+ a_contiguous.const_data_ptr<scalar_t>(),
66
+ output.mutable_data_ptr<scalar_t>()); });
67
+ }
68
+ else
69
+ {
70
+ AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
71
+ x.scalar_type(), "host_lti_shared_linear_recurrence", [&]
72
+ { host_lti_shared_linear_recurrence<scalar_t>(
73
+ n_batches, n_steps,
74
+ a_contiguous.item<scalar_t>(),
75
+ output.mutable_data_ptr<scalar_t>()); });
76
+ }
77
+
78
+ return output.slice(1, 1, output.size(1))
79
+ .contiguous(); // Remove the initial state from the output
80
+ }
81
+
82
+ TORCH_LIBRARY_IMPL(philtorch, CPU, m)
83
+ {
84
+ m.impl("lti_recur", &lti_recur_cpu_impl);
85
+ }
@@ -0,0 +1,114 @@
1
+ #include <torch/script.h>
2
+ #include <torch/torch.h>
3
+
4
+ template <typename scalar_t>
5
+ void host_lti_batch_mat_recur_N_order(int B, int T, int N,
6
+ const scalar_t *A, // [B][N][N]
7
+ scalar_t *H_out // [B][T][N]
8
+ )
9
+ {
10
+ // Process each batch in parallel
11
+ #pragma omp parallel for
12
+ for (int b = 0; b < B; ++b)
13
+ {
14
+ const scalar_t *A_b = A + b * N * N;
15
+ scalar_t *H_b = H_out + b * T * N;
16
+
17
+ // t loop
18
+ for (int t = 1; t < T; ++t)
19
+ {
20
+ const scalar_t *H_prev = H_b + (t - 1) * N;
21
+ scalar_t *H_curr = H_b + t * N;
22
+
23
+ for (int i = 0; i < N; ++i)
24
+ {
25
+ scalar_t sum = 0;
26
+ const scalar_t *Arow = A_b + i * N;
27
+ for (int j = 0; j < N; ++j)
28
+ sum += Arow[j] * H_prev[j];
29
+ H_curr[i] += sum;
30
+ }
31
+ }
32
+ }
33
+ }
34
+
35
+ template <typename scalar_t>
36
+ void host_lti_share_mat_recur_N_order(int B, int T, int N,
37
+ const scalar_t *A, // N][N]
38
+ scalar_t *H_out // [B][T][N]
39
+ )
40
+ {
41
+ // Process each batch in parallel
42
+ #pragma omp parallel for
43
+ for (int b = 0; b < B; ++b)
44
+ {
45
+ scalar_t *H_b = H_out + b * T * N;
46
+ // t loop
47
+ for (int t = 1; t < T; ++t)
48
+ {
49
+ const scalar_t *H_prev = H_b + (t - 1) * N;
50
+ scalar_t *H_curr = H_b + t * N;
51
+ for (int i = 0; i < N; ++i)
52
+ {
53
+ scalar_t sum = 0;
54
+ const scalar_t *Arow = A + i * N;
55
+ for (int j = 0; j < N; ++j)
56
+ sum += Arow[j] * H_prev[j];
57
+ H_curr[i] += sum;
58
+ }
59
+ }
60
+ }
61
+ }
62
+
63
+ at::Tensor lti_mat_recur_N_order_cpu_impl(const at::Tensor &A,
64
+ const at::Tensor &zi,
65
+ const at::Tensor &x)
66
+ {
67
+ TORCH_CHECK(zi.scalar_type() == x.scalar_type(),
68
+ "zi must have the same scalar type as input");
69
+ TORCH_CHECK(A.scalar_type() == x.scalar_type(),
70
+ "A must have the same scalar type as input");
71
+ TORCH_CHECK(A.dim() == 2 || A.dim() == 3, "A must be a 2D or 3D tensor");
72
+ TORCH_CHECK(x.size(2) == A.size(-1),
73
+ "Last dimension of x must match last dimension of A");
74
+ TORCH_CHECK(A.size(-1) == A.size(-2),
75
+ "Last two dimensions of A must be equal");
76
+
77
+ auto n_steps = x.size(1) + 1; // +1 for the initial state
78
+ auto n_batches = x.size(0);
79
+ auto order = A.size(-1);
80
+
81
+ auto A_contiguous = A.contiguous();
82
+ auto out = at::cat({zi.unsqueeze(1), x}, 1).contiguous();
83
+
84
+ if (A.dim() == 3)
85
+ {
86
+ AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1(
87
+ at::kLong, x.scalar_type(), "host_lti_batch_mat_recur_N_order",
88
+ [&]
89
+ {
90
+ host_lti_batch_mat_recur_N_order<scalar_t>(
91
+ n_batches, n_steps, order,
92
+ A_contiguous.const_data_ptr<scalar_t>(),
93
+ out.mutable_data_ptr<scalar_t>());
94
+ });
95
+ }
96
+ else
97
+ {
98
+ // Shared
99
+ AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1(
100
+ at::kLong, x.scalar_type(), "host_lti_share_mat_recur_N_order", [&]
101
+ { host_lti_share_mat_recur_N_order<scalar_t>(
102
+ n_batches, n_steps, order,
103
+ A_contiguous.const_data_ptr<scalar_t>(),
104
+ out.mutable_data_ptr<scalar_t>()); });
105
+ }
106
+ return out.slice(1, 1, n_steps)
107
+ .contiguous(); // Remove the initial state from the output
108
+ }
109
+
110
+ TORCH_LIBRARY_IMPL(philtorch, CPU, m)
111
+ {
112
+ m.impl("lti_recurN", &lti_mat_recur_N_order_cpu_impl);
113
+ m.impl("lti_recur2", &lti_mat_recur_N_order_cpu_impl);
114
+ }