MilesCranmer / symbolic_deep_learning

Code for "Discovering Symbolic Models from Deep Learning with Inductive Biases"

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Support for Windows systems that cannot install jax?

jstan15 opened this issue · comments

Hi, thank you for your amazing work on this and PySR and for making it available for us to use.

I have encountered a problem with the demo file because it requires jaxlib, which isn't available on windows. Additionally, I don't have a gpu either. Based on my understanding, jax is only being used to generate the data. However I proficient enough to be able to understand the code and write a general version that works with other types of data so that jaxlib isn't required, without first being able to run the entire demo to understand each step in detail

Would you please advice me on what I can do to make this work on windows?
Thank you very much.
Regards,
jstan

Unfortunately you can't run the data-generating parts of the example without having jax. However, jax is only used to generate the simulation data (as you pointed out!), so you could still do the other parts if you cut out that code. What I would do would be: run the simulation in Google colab (see link on the readme). Download the data, then on your windows machine, load the data that you downloaded instead of generating it. JAX isn't used anywhere else so the rest should be okay. Then just remove the .cuda() parts to have it work without a GPU.

In general it would be easiest to use the colab version though!

You could also try to get jax set up on windows using the advice in this thread: google/jax#438 (comment), but I'm assuming it will take a bit of work.

Cheers,
Miles

I see, Thank you for your reply.