feat: expose linear recurrence function and update __all__ exports

Author: yoyolicoris

torchlpc/__init__.py

@@ -23,11 +23,41 @@
     warnings.warn("Custom extension not loaded. Falling back to Numba implementation.")
 
 from .core import LPC
-
-# from .parallel_scan import WARPSIZE
 from .recurrence import Recurrence
 
-__all__ = ["sample_wise_lpc"]
+__all__ = ["sample_wise_lpc", "linear_recurrence"]
+
+
+def linear_recurrence(
+    a: torch.Tensor, x: torch.Tensor, zi: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Compute linear recurrence using the recurrence relation:
+    y[n] = x[n] + a[n] * (x[n-1] + a[n-1] * (x[n-2] + a[n-2] * (...(x[0] + a[0] * zi)...)))
+
+    Args:
+        a (torch.Tensor): Coefficients of the recurrence relation.
+        x (torch.Tensor): Input signal.
+        zi (torch.Tensor, optional): Initial conditions. Defaults to zero if not provided.
+
+    Shape:
+        - a: :math:`(B, T)`
+        - x: :math:`(B, T)`
+        - zi: :math:`(B,)`
+
+    Returns:
+        Output signal with the same shape as x.
+    """
+    assert a.shape == x.shape
+    assert a.ndim == 2
+    assert x.ndim == 2
+    B, _ = a.shape
+
+    if zi is None:
+        zi = a.new_zeros(B)
+    else:
+        assert zi.shape == (B,)
+
+    return Recurrence.apply(a, x, zi)  # type: ignore
 
 
 def sample_wise_lpc(
@@ -64,10 +94,8 @@ def sample_wise_lpc(
     else:
         assert zi.shape == (B, order)
 
-    # if order == 1 and x.is_cuda and B * WARPSIZE < T:
-    #     return RecurrenceCUDA.apply(-a.squeeze(2), x, zi.squeeze(1))
     if order == 1:
-        y = Recurrence.apply(-a.squeeze(2), x, zi.squeeze(1))
+        y = linear_recurrence(-a.squeeze(2), x, zi.squeeze(1))
     else:
         y = LPC.apply(x, a, zi)