Nice work @ordabayevy! The algorithm looks good to me.

Could you update the docstring to reflect complexity changes, in particular the two paragraphs

This uses [1] to parallelize over time, achieving O(log(time)) parallel
complexity for computing :meth:log_prob and :meth:filter.

and

The :meth:sample method is sequential (not parallized), slow, and memory
inefficient. It is intended for data generation only and is not recommended
during inference.