PyTorch only cauchy kernel for easier test

[bratao]

Hello,

I loved Structured State Spaces, and obteined a fantastic performance compared to LSTM/SRU/Transformers.

I want to introduce S4 to some researchers and students, and the self-contained S4 layer is super great! However it requires the "cautchy kernel".

There are two versions of cautchy kernel, a cuda and a Pykeops version.

However extensions/cauchy requires cuda and Pykeops do not support Windows, and the target people have a very diverse number of environments.

Would be possible to HazyResearch team to implement a self-contained S4 layer including a simpler pytorch cautchy kernel ?

[albertfgu]

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

[bratao]

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

[albertfgu]

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/

[albertfgu]

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

[bratao]

Works perfectly @albertfgu , thank you!