Possibility to speed up further with JAX?

[SimonBiggs]

Have you seen Google's JAX project?

https://github.com/google/jax

Do you think it might feasible to speed this up further with that?

[rafael-fuente]

Thanks for the suggestion!
Some people suggested me to use JAX, cupy, pytorch etc instead of numpy to speed up calculations with a GPU. I tested using GPU backend pytorch tensors, but the main problem I see is that although it runs faster for single calls, for multiple calls kills the performance due the function call overhead.
I read that JAX can compile several functions using a jit compiler, so it is likely that some parts can be optimized. But the main problem I'm going to face is that my source code has a lot recursion and branching because it's made to be as readable as possible. But definitely when I have time I have to give it a try.

[SimonBiggs]

But the main problem I'm going to face is that my source code has a lot recursion and branching

Yup, that makes sense.

[sparthir]

So this means that multiprocessing is also complicated to implement?

I've seen it done here: https://github.com/NMVHS/PyTracer

[rafael-fuente]

So this means that multiprocessing is also complicated to implement?

Multiprocessing works on CPU cores so it doesn't have these GPU limitations. What happens it's that my raytracer already uses multiple cores because lot of Numpy methods uses multithreading so it cannot be made faster using multiprocessing.

About the PyTracer project I already have talked with the author and made some tests. It's also a cool project.
But a pure Python implementation with multiprocessing isn't faster than a Numpy implementation on an usual PC. (~4 cores)
The reason is because Python loops are very slow even running on multiple cores. So throwing the pixel loop to Numpy and make use of their multithreaded compiled methods is the best idea I've seen.