PyTorch only cauchy kernel for easier test about s4 HOT 5 CLOSED

These extensions are required because the algorithm requires functionality that is not supported in vanilla PyTorch. There is a pure PyTorch implementation of the kernel that is slower and uses more memory. To use it, in the line https://github.com/HazyResearch/state-spaces/blob/2af126108991d214fd82ffc1899f6d4e31a1eda3/src/models/sequence/ss/kernel.py#L322 replace cauchy.cauchy_conj with cauchy.cauchy_conj_slow.
Or if you're using the standalone module, replace the cauchy function with the one found here https://github.com/HazyResearch/state-spaces/blob/2af126108991d214fd82ffc1899f6d4e31a1eda3/src/models/functional/cauchy.py#L124

from s4.

That worked @albertfgu , thank you!
I will open a PR to use this as fallback on the standalone module

from s4.

No problem. I can also make the requisite changes to fall back on this - should be pretty simple.

Do keep in mind again that this version is less efficient, so it might be useful for preliminary testing but we would recommend trying to set up the full environment if possible.

Also, there is an excellent port of S4 to JAX here if that might be easier to set up: https://srush.github.io/annotated-s4/

from s4.

I just added a fallback to the slow version. Can you test this?

from s4.

Works perfectly @albertfgu , thank you!

from s4.