Getting `KeyError: 'nvrtc'` on CPU-only machine

[geoffreyangus]

Versions:

• Python 3.9.14
• macOS 12.6
• state-spaces@292984c
• ludwig@1d8154f

Hello!

I'd like to add S4 into the Ludwig OSS project. I've successfully imported and initialized the S4 module. However, in the forward pass, I am running into the following error:

[KeOps] Warning : Cuda libraries were not detected on the system ; using cpu only mode
Traceback (most recent call last):
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 1670, in <module>
    main()
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 1666, in main
    module(torch.randn(2, 16, 100))
  File "/Users/geoffreyangus/repositories/predibase/ludwig/venv39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 1582, in forward
    k, k_state = self.kernel(L=L_kernel, rate=rate, state=state)  # (C H L) (B C H L)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/venv39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 1388, in forward
    return self.kernel(state=state, L=L, rate=rate)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/venv39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 821, in forward
    r = cauchy_conj(v, z, w)
  File "/Users/geoffreyangus/repositories/predibase/ludwig/ludwig/modules/s4_modules.py", line 59, in cauchy_conj
    r = 2 * cauchy_mult(v, z, w, backend="GPU")
  File "/Users/geoffreyangus/repositories/predibase/ludwig/venv39/lib/python3.9/site-packages/pykeops/torch/generic/generic_red.py", line 624, in __call__
    out = GenredAutograd.apply(
  File "/Users/geoffreyangus/repositories/predibase/ludwig/venv39/lib/python3.9/site-packages/pykeops/torch/generic/generic_red.py", line 78, in forward
    myconv = keops_binder["nvrtc" if tagCPUGPU else "cpp"](
KeyError: 'nvrtc'

Here is the code snippet I am currently running from within the file implementing the S4 class:

def main():
    module = S4(
        16,
        gate=4,  # Multiplicative gating layer that also expands dimension by factor of 4
        bottleneck=4,  # Reduce dimension of SSM by factor of 4
        measure="legs",  # Randomly initialize A
        dt_min=1.0,
        dt_max=1.0,  # Initialize dt to 1.0
        lr={"dt": 0.0, "B": 0.0},  # Freeze B and dt
    )
    module(torch.randn(2, 16, 100))


if __name__ == "__main__":
    main()

To reproduce the error:

1. Download Ludwig on a CPU-only machine. cd into the repository.
2. Checkout the hack-s4 branch.
3. Run python ludwig/modules/s4_modules.py.

Let me know what you think, thanks!