JAX based implementation

[isty2e]

JAX provides even more speedup over CuPy, with additional nice functionalities and transformations, and it introduces less problems while configuring relevant environments. The migration will be quite time consuming, of course, but I wonder if such a process can be considered viable.

[jlparkI]

Switching to another framework (e.g. JAX, PyTorch) is definitely an interesting possibility, but it would need to overcome some challenges.

The main benefit to JAX would be JIT, which would provide a speedup over Cupy for some operations. It would also simplify the code in some places. We would need however to figure out how to incorporate the C / Cuda code. Out of necessity, some of the random feature generation code has to be written in C and Cuda (to apply the fast Hadamard transform), and we've gradually been rewriting all of the steps involved in random feature generation for different kernels in C / Cuda to reduce overhead and improve speed. (Random feature generation for the RBF / Matern kernels will be about 20-25% faster on GPU in the next release, for example, and we're working on accelerating the graph and sequence kernels as well). I am not sure how well JAX JIT plays with C / Cuda extensions. It's undoubtedly possible, but there are a number of issues that would need to be figured out.