Skip to content

Time-Domain Filtering

Time-domain filtering functions that run efficiently on CPU and GPU, and are differentiable in all arguments. The implementations are strongly based on philtorch and torchlpc.

Time-Invariant Filtering

Functions that apply linear time-invariant (LTI) filters to time-domain signals. Filter coefficients are constant over time.

korvax.filter.lti.lfilter 

lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_a"] | None = None,
    b: Float[Array, " n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[False] = False,
    transposed: bool = True,
) -> Float[Array, " n_samples"]

lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_a"] | None = None,
    b: Float[Array, " n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[True],
    transposed: bool = True,
) -> tuple[
    Float[Array, " n_samples"], Float[Array, " order"]
]

lfilter(
    x,
    a=None,
    b=None,
    zi=None,
    *,
    return_zi=False,
    transposed=True,
)

Apply a time-invariant filter to the input signal.

Filtering is implemented using the state-space implementations with parallel associative scans as described in [1]. In the time-invariant case, this is also efficient at higher filter orders.

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
a Float[Array, ' n_a'] | None

Denominator (IIR) coefficients a_1, a_2, ... of shape (n_a). a_0 = 1 is implied. Can be None for all-zero filters.

 None
b Float[Array, ' n_b'] | None

Numerator (FIR) coefficients b_0, b_1, ... of shape (n_b). Can be None for all-pole filters.

 None
zi Float[Array, ' order'] | None

Initial conditions of shape (order,), where order=max(n_a, n_b) - 1. If None, zeros are used.

 None
return_zi bool

If True, return the final conditions along with the output.

 False
transposed bool

Whether to use transposed direct form II structure (default). Uses direct form II otherwise.

 True

Returns:

Type Description
tuple[Float[Array, ' n_samples'], Float[Array, ' order']] | Float[Array, ' n_samples']

If return_zi is False, returns the filtered signal of shape (n_samples,). If return_zi is True, returns a tuple containing:

  • Filtered signal of shape (n_samples,)
  • Final conditions of shape (order,), where order=max(n_a, n_b) - 1.
References

[1] C.-Y. Yu and G. Fazekas. "Accelerating Automatic Differentiation of Direct Form Digital Filters", DiffSys Workshap at EurIPS, 2025.

Source code in src/korvax/filter/lti.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_a"] | None = None,
    b: Float[Array, " n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: bool = False,
    transposed: bool = True,
) -> (
    tuple[Float[Array, " n_samples"], Float[Array, " order"]]
    | Float[Array, " n_samples"]
):
    """Apply a time-invariant filter to the input signal.

    Filtering is implemented using the state-space implementations with parallel associative scans as described in [1]. In the time-invariant case, this is also efficient at higher filter orders.

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.


    Args:
        x: Input signal of shape `(n_samples,)`.
        a: Denominator (IIR) coefficients `a_1, a_2, ...` of shape `(n_a)`. `a_0 = 1` is implied. Can be `None` for all-zero filters.
        b: Numerator (FIR) coefficients `b_0, b_1, ...` of shape `(n_b)`. Can be `None` for all-pole filters.
        zi: Initial conditions of shape `(order,)`, where `order=max(n_a, n_b) - 1`. If `None`, zeros are used.
        return_zi: If `True`, return the final conditions along with the output.
        transposed: Whether to use transposed direct form II structure (default). Uses direct form II otherwise.

    Returns:
        If `return_zi` is `False`, returns the filtered signal of shape `(n_samples,)`. If `return_zi` is `True`, returns a tuple containing:

            - Filtered signal of shape `(n_samples,)`
            - Final conditions of shape `(order,)`, where `order=max(n_a, n_b) - 1`.

    References:
        [1] C.-Y. Yu and G. Fazekas. "Accelerating Automatic Differentiation of Direct Form Digital Filters", DiffSys Workshap at EurIPS, 2025.
    """
    if a is None:
        a = jnp.array([], dtype=x.dtype)

    if b is None:
        b = jnp.array([1.0], dtype=x.dtype)

    order = max(a.shape[-1], b.shape[-1] - 1)

    if zi is None:
        zi = jnp.zeros((order,), dtype=x.dtype)
    else:
        assert zi.shape == (order,)

    assert zi is not None

    if b.shape[-1] < order + 1:
        b = jnp.pad(b, (0, order + 1 - b.shape[-1]))
    if a.shape[-1] < order:
        a = jnp.pad(a, (0, order - a.shape[-1]))

    A = _companion(a)
    b0 = b[:1]
    C = b[1:] - a * b0
    D = b[0]

    if transposed:
        A = A.T
        z = x[:, None] * C

    else:
        z = jnp.pad(x[:, None], ((0, 0), (0, order - 1)))

    v = _recurrence(A, jnp.concatenate([zi[None, :], z], axis=0))

    if transposed:
        y = v[:-1, 0] + D * x
    else:
        y = jnp.dot(v[:-1, :], C) + D * x

    if return_zi:
        return y, v[-1, :]
    else:
        return y

korvax.filter.lti.sosfilt 

sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections 2"],
    b: Float[Array, " n_sections 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: Literal[False] = False,
) -> Float[Array, " n_samples"]

sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections 2"],
    b: Float[Array, " n_sections 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: Literal[True],
) -> tuple[
    Float[Array, " n_samples"],
    Float[Array, " n_sections 2"],
]

sosfilt(x, a, b, zi=None, *, return_zi=False)

Apply a cascade of time-invariant second-order filters (biquads) to the input signal.

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
a Float[Array, ' n_sections 2']

Denominator (IIR) coefficients of shape (n_sections, 2).

 required
b Float[Array, ' n_sections 3']

Numerator (FIR) coefficients of shape (n_sections, 3).

 required
zi Float[Array, ' n_sections 2'] | None

Initial conditions of shape (n_sections, 2). If None, zeros are used.

 None
return_zi bool

If True, return the final conditions along with the output.

 False

Returns:

Type Description
tuple[Float[Array, ' n_samples'], Float[Array, ' n_sections 2']] | Float[Array, ' n_samples']

If return_zi is False, returns the filtered signal of shape (n_samples,). If return_zi is True, returns a tuple containing: - Filtered signal of shape (n_samples,) - Final conditions of shape (n_sections, 2).
Source code in src/korvax/filter/lti.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
def sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections 2"],
    b: Float[Array, " n_sections 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: bool = False,
) -> (
    tuple[Float[Array, " n_samples"], Float[Array, " n_sections 2"]]
    | Float[Array, " n_samples"]
):
    """Apply a cascade of time-invariant second-order filters (biquads) to the input signal.

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.


    Args:
        x: Input signal of shape `(n_samples,)`.
        a: Denominator (IIR) coefficients of shape `(n_sections, 2)`.
        b: Numerator (FIR) coefficients of shape `(n_sections, 3)`.
        zi: Initial conditions of shape `(n_sections, 2)`. If `None`, zeros are used.
        return_zi: If `True`, return the final conditions along with the output.

    Returns:
        If `return_zi` is `False`, returns the filtered signal of shape `(n_samples,)`. If `return_zi` is `True`, returns a tuple containing:
            - Filtered signal of shape `(n_samples,)`
            - Final conditions of shape `(n_sections, 2)`.
    """
    n_sections = a.shape[0]

    if zi is None:
        zi = jnp.zeros((n_sections, 2), dtype=x.dtype)
    else:
        assert zi.shape == (n_sections, 2)

    def _section(carry, inputs):
        a_, b_, zi_ = inputs
        return lfilter(carry, a_, b_, zi=zi_, return_zi=True)

    y, zi_out = lax.scan(_section, x, (a, b, zi))
    if return_zi:
        return y, zi_out
    else:
        return y

Time-Varying Filtering

Functions that apply linear time-varying (LTV) filters to time-domain signals. Filter coefficients can change at audio sample rate.

korvax.filter.ltv.lfilter 

lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples n_a"] | None = None,
    b: Float[Array, " n_samples n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[False] = False,
    transposed: bool = False,
) -> Float[Array, " n_samples"]

lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples n_a"] | None = None,
    b: Float[Array, " n_samples n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[True],
    transposed: bool = False,
) -> tuple[
    Float[Array, " n_samples"], Float[Array, " order"]
]

lfilter(
    x,
    a=None,
    b=None,
    zi=None,
    *,
    return_zi=False,
    transposed=False,
)

Apply a linear filter with time-varying coefficients to the input signal.

Filtering is implemented using the state-space implementations with parallel associative scans as described in [1]. No diagonalization is implemented currently! For time-varying filters with order > ~4, combining ltv.fir and ltv.allpole will likely be a lot more performant.

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
a Float[Array, ' n_samples n_a'] | None

Time-varying denominator (IIR) coefficients a_1, a_2, ... of shape (n_samples, n_a). a_0 = 1 is implied. Can be None for all-zero filters.

 None
b Float[Array, ' n_samples n_b'] | None

Time-varying numerator (FIR) coefficients b_0, b_1, ... of shape (n_samples, n_b). Can be None for all-pole filters.

 None
zi Float[Array, ' order'] | None

Initial conditions of shape (order,), where order=max(n_a, n_b - 1). If None, zeros are used.

 None
return_zi bool

If True, return the final conditions along with the output.

 False
transposed bool

Whether to use transposed direct form II structure. Uses direct form II if False (default).

 False

Returns:

Type Description
tuple[Float[Array, ' n_samples'], Float[Array, ' order']] | Float[Array, ' n_samples']

If return_zi is False, returns the filtered signal of shape (n_samples,). If return_zi is True, returns a tuple containing:

  • Filtered signal of shape (n_samples,)
  • Final conditions of shape (order,), where order=max(n_a, n_b) - 1.
References

[1] C.-Y. Yu and G. Fazekas. "Accelerating Automatic Differentiation of Direct Form Digital Filters", DiffSys Workshap at EurIPS, 2025.

Source code in src/korvax/filter/ltv.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def lfilter(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples n_a"] | None = None,
    b: Float[Array, " n_samples n_b"] | None = None,
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: bool = False,
    transposed: bool = False,
) -> (
    tuple[Float[Array, " n_samples"], Float[Array, " order"]]
    | Float[Array, " n_samples"]
):
    """Apply a linear filter with time-varying coefficients to the input signal.

    Filtering is implemented using the state-space implementations with parallel associative scans as described in [1].
    No diagonalization is implemented currently! For time-varying filters with order > ~4, combining [`ltv.fir`][korvax.filter.ltv.fir] and [`ltv.allpole`][korvax.filter.ltv.allpole] will likely be a lot more performant.

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.


    Args:
        x: Input signal of shape `(n_samples,)`.
        a: Time-varying denominator (IIR) coefficients `a_1, a_2, ...` of shape `(n_samples, n_a)`. `a_0 = 1` is implied. Can be `None` for all-zero filters.
        b: Time-varying numerator (FIR) coefficients `b_0, b_1, ...` of shape `(n_samples, n_b)`. Can be `None` for all-pole filters.
        zi: Initial conditions of shape `(order,)`, where `order=max(n_a, n_b - 1)`. If `None`, zeros are used.
        return_zi: If `True`, return the final conditions along with the output.
        transposed: Whether to use transposed direct form II structure. Uses direct form II if False (default).

    Returns:
        If `return_zi` is `False`, returns the filtered signal of shape `(n_samples,)`. If `return_zi` is `True`, returns a tuple containing:

            - Filtered signal of shape `(n_samples,)`
            - Final conditions of shape `(order,)`, where `order=max(n_a, n_b) - 1`.

    References:
        [1] C.-Y. Yu and G. Fazekas. "Accelerating Automatic Differentiation of Direct Form Digital Filters", DiffSys Workshap at EurIPS, 2025.
    """
    n_samples = x.shape[0]

    if a is None:
        a = jnp.empty((n_samples, 0), dtype=x.dtype)

    if b is None:
        b = jnp.ones((n_samples, 1), dtype=x.dtype)

    order = max(a.shape[-1], b.shape[-1] - 1)

    if zi is None:
        zi = jnp.zeros((order,), dtype=x.dtype)
    else:
        assert zi.shape == (order,)

    assert zi is not None

    if b.shape[-1] < order + 1:
        b = jnp.pad(b, ((0, 0), (0, order + 1 - b.shape[-1])))
    if a.shape[-1] < order:
        a = jnp.pad(a, ((0, 0), (0, order - a.shape[-1])))

    A = _companion(a)
    b0 = b[:, :1]
    C = b[:, 1:] - a * b0
    D = b[:, 0]

    if transposed:
        A = A.mT
        z = x[:, None] * C

    else:
        z = jnp.pad(x[:, None], ((0, 0), (0, order - 1)))

    v = _recurrence(A, z, zi)

    if transposed:
        y = v[:-1, 0] + D * x
    else:
        y = jnp.linalg.vecdot(v[:-1, :], C) + D * x

    if return_zi:
        return y, v[-1, :]
    else:
        return y

korvax.filter.ltv.sosfilt 

sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections n_samples 2"],
    b: Float[Array, " n_sections n_samples 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: Literal[False] = False,
) -> Float[Array, " n_samples"]

sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections n_samples 2"],
    b: Float[Array, " n_sections n_samples 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: Literal[True],
) -> tuple[
    Float[Array, " n_samples"],
    Float[Array, " n_sections 2"],
]

sosfilt(x, a, b, zi=None, *, return_zi=False)

Apply a cascade of second-order filters (biquads) with time-varying coefficients to the input signal.

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
a Float[Array, ' n_sections n_samples 2']

Denominator (IIR) coefficients of shape (n_sections, n_samples, 2).

 required
b Float[Array, ' n_sections n_samples 3']

Numerator (FIR) coefficients of shape (n_sections, n_samples, 3).

 required
zi Float[Array, ' n_sections 2'] | None

Initial conditions of shape (n_sections, 2). If None, zeros are used.

 None
return_zi bool

If True, return the final conditions along with the output.

 False

Returns:

Type Description
tuple[Float[Array, ' n_samples'], Float[Array, ' n_sections 2']] | Float[Array, ' n_samples']

If return_zi is False, returns the filtered signal of shape (n_samples,). If return_zi is True, returns a tuple containing: - Filtered signal of shape (n_samples,) - Final conditions of shape (n_sections, 2).
Source code in src/korvax/filter/ltv.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
def sosfilt(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_sections n_samples 2"],
    b: Float[Array, " n_sections n_samples 3"],
    zi: Float[Array, " n_sections 2"] | None = None,
    *,
    return_zi: bool = False,
) -> (
    tuple[Float[Array, " n_samples"], Float[Array, " n_sections 2"]]
    | Float[Array, " n_samples"]
):
    """Apply a cascade of second-order filters (biquads) with time-varying coefficients to the input signal.

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.


    Args:
        x: Input signal of shape `(n_samples,)`.
        a: Denominator (IIR) coefficients of shape `(n_sections, n_samples, 2)`.
        b: Numerator (FIR) coefficients of shape `(n_sections, n_samples, 3)`.
        zi: Initial conditions of shape `(n_sections, 2)`. If `None`, zeros are used.
        return_zi: If `True`, return the final conditions along with the output.

    Returns:
        If `return_zi` is `False`, returns the filtered signal of shape `(n_samples,)`. If `return_zi` is `True`, returns a tuple containing:
            - Filtered signal of shape `(n_samples,)`
            - Final conditions of shape `(n_sections, 2)`.
    """
    n_sections = a.shape[0]

    if zi is None:
        zi = jnp.zeros((n_sections, 2), dtype=x.dtype)
    else:
        assert zi.shape == (n_sections, 2)

    def _section(carry, inputs):
        a_, b_, zi_ = inputs
        return lfilter(carry, a_, b_, zi=zi_, return_zi=True)

    y, zi_out = lax.scan(_section, x, (a, b, zi))
    if return_zi:
        return y, zi_out
    else:
        return y

korvax.filter.ltv.fir 

fir(x, b, zi=None)

Apply a linear FIR filter with time-varying coefficients to the input signal.

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
b Float[Array, ' n_samples n_b']

Time-varying FIR coefficients b_0, b_1, ... of shape (n_samples, order+1).

 required
zi Float[Array, ' n_b-1'] | None

Initial conditions of shape (order,). If None, zeros are used.

 None

Returns:

Type Description
Float[Array, ' n_samples']

Filtered signal of shape (n_samples,).
Source code in src/korvax/filter/ltv.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
def fir(
    x: Float[Array, " n_samples"],
    b: Float[Array, " n_samples n_b"],
    zi: Float[Array, " n_b-1"] | None = None,
) -> Float[Array, " n_samples"]:
    """Apply a linear FIR filter with time-varying coefficients to the input signal.

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.

    Args:
        x: Input signal of shape `(n_samples,)`.
        b: Time-varying FIR coefficients `b_0, b_1, ...` of shape `(n_samples, order+1)`.
        zi: Initial conditions of shape `(order,)`. If `None`, zeros are used.

    Returns:
        Filtered signal of shape `(n_samples,)`.
    """
    order = b.shape[-1] - 1
    if zi is None:
        zi = jnp.zeros((order,), dtype=x.dtype)

    n_samples = x.shape[0]
    x = jnp.r_[zi, x]

    frames = jax.vmap(
        partial(lax.dynamic_slice_in_dim, operand=x, slice_size=order + 1)
    )(start_index=jnp.arange(n_samples))

    frames = jnp.flip(frames, axis=1)

    return jnp.linalg.vecdot(frames, b)

korvax.filter.ltv.allpole 

allpole(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples order"],
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[False] = False,
) -> Float[Array, " n_samples"]

allpole(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples order"],
    zi: Float[Array, " order"] | None = None,
    *,
    return_zi: Literal[True],
) -> tuple[
    Float[Array, " n_samples"], Float[Array, " order"]
]

allpole(x, a, zi=None, return_zi=False)

Apply a time-varying all-pole filter to the input signal.

Port of torchlpc.sample_wise_lpc. Uses the efficient differentiation method proposed in [1].

This function only operates on 1D signals, use jax.vmap to apply it to batched inputs.

Parameters:

Name Type Description Default
x Float[Array, ' n_samples']

Input signal of shape (n_samples,).

 required
a Float[Array, ' n_samples order']

Time-varying all-pole coefficients of shape (n_samples, order).

 required
zi Float[Array, ' order'] | None

Initial conditions of shape (order,). If None, zeros are used.

 None
return_zi bool

If True, return the final conditions along with the output.

 False

Returns:

Type Description
Float[Array, ' n_samples'] | tuple[Float[Array, ' n_samples'], Float[Array, ' order']]

If return_zi is False, returns the filtered signal of shape (n_samples,). If return_zi is True, returns a tuple containing:

  • Filtered signal of shape (n_samples,)
  • Final conditions of shape (order,)
References

[1] C.-Y. Yu, C. Mitcheltree, A. Carson, S. Bilbao, J. D. Reiss, and G. Fazekas. "Differentiable All-Pole Filters for Time-Varying Audio Systems," in Proc. DAFx, 2024.

Source code in src/korvax/filter/ltv.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
@partial(jax.custom_vjp, nondiff_argnames=("return_zi",))
def allpole(
    x: Float[Array, " n_samples"],
    a: Float[Array, " n_samples order"],
    zi: Float[Array, " order"] | None = None,
    return_zi: bool = False,
) -> (
    Float[Array, " n_samples"]
    | tuple[Float[Array, " n_samples"], Float[Array, " order"]]
):
    """Apply a time-varying all-pole filter to the input signal.

    Port of `torchlpc.sample_wise_lpc`. Uses the efficient differentiation method proposed in [1].

    This function only operates on 1D signals, use `jax.vmap` to apply it to batched inputs.


    Args:
        x: Input signal of shape `(n_samples,)`.
        a: Time-varying all-pole coefficients of shape `(n_samples, order)`.
        zi: Initial conditions of shape `(order,)`. If `None`, zeros are used.
        return_zi: If `True`, return the final conditions along with the output.

    Returns:
        If `return_zi` is `False`, returns the filtered signal of shape `(n_samples,)`. If `return_zi` is `True`, returns a tuple containing:

            - Filtered signal of shape `(n_samples,)`
            - Final conditions of shape `(order,)`

    References:
        [1] C.-Y. Yu, C. Mitcheltree, A. Carson, S. Bilbao, J. D. Reiss, and G. Fazekas. "Differentiable All-Pole Filters for Time-Varying Audio Systems," in Proc. DAFx, 2024.
    """

    order = a.shape[-1]
    if zi is None:
        zi = jnp.zeros((order,), dtype=x.dtype)
    x = jnp.r_[zi, x]

    def _call(target: str):
        return jax.ffi.ffi_call(
            target, jax.ShapeDtypeStruct(x.shape, x.dtype), vmap_method="broadcast_all"
        )

    out = lax.platform_dependent(
        x,
        a,
        default=_call("allpole_cpu"),
        cpu=_call("allpole_cpu"),
        cuda=_call("allpole_cuda"),
    )

    if return_zi:
        return out[..., order:], out[..., -order:]
    return out[..., order:]