forward mode optimisation

* merge two filtering into one

* update version number

CITATION.cff

@@ -1,10 +1,10 @@
 cff-version: 1.2.0
 message: "If you use this software, please cite it as below."
 authors:
-- family-names: "Yu"
-  given-names: "Chin-Yun"
-  orcid: "https://orcid.org/0000-0003-3782-2713"
+  - family-names: "Yu"
+    given-names: "Chin-Yun"
+    orcid: "https://orcid.org/0000-0003-3782-2713"
 title: "TorchLPC: fast, efficient, and differentiable time-varying LPC filtering in PyTorch"
-version: 0.1.0
+version: 0.3.1
 date-released: 2023-07-09
 url: "https://github.com/yoyololicon/torchlpc"

setup.py

@@ -1,7 +1,7 @@
 import setuptools
 
 NAME = "torchlpc"
-VERSION = "0.3"
+VERSION = "0.3.1"
 MAINTAINER = "Chin-Yun Yu"
 EMAIL = "[email protected]"
 

torchlpc/core.py

@@ -189,18 +189,14 @@ def jvp(
         A, y, zi = ctx.A, ctx.y, ctx.zi
         *_, order = A.shape
 
-        grad_y = 0
-
-        if grad_x is not None:
-            grad_y_from_x_zi = LPC.apply(grad_x, A, grad_zi)
-            grad_y = grad_y_from_x_zi
+        fwd_zi = grad_zi if grad_zi is not None else torch.zeros_like(zi)
+        fwd_x = grad_x if grad_x is not None else torch.zeros_like(y)
 
         if grad_A is not None:
             unfolded_y = (
                 torch.cat([zi.flip(1), y[:, :-1]], dim=1).unfold(1, order, 1).flip(2)
             )
-            grad_A_input = -torch.sum(unfolded_y * grad_A, dim=2)
-            grad_y_from_A = LPC.apply(grad_A_input, A, torch.zeros_like(zi))
-            grad_y = grad_y + grad_y_from_A
+            fwd_A = -torch.sum(unfolded_y * grad_A, dim=2)
+            fwd_x = fwd_x + fwd_A
 
-        return grad_y
+        return LPC.apply(fwd_x, A, fwd_zi)