Run Error, need help! about fbpinns HOT 8 CLOSED

I faced the same issue. I fixed it by installing pip install -U "jax[cpu]" link, and rest of the dependencies
"optax>=0.1.4", "numpy>=1.24.2", "scipy>=1.10.1", "matplotlib>=3.7.1", "tensorboardX>=2.6", "ipython>=8.12.0"]

from fbpinns.

I faced the same issue. I fixed it by installing pip install -U "jax[cpu]" link, and rest of the dependencies "optax>=0.1.4", "numpy>=1.24.2", "scipy>=1.10.1", "matplotlib>=3.7.1", "tensorboardX>=2.6", "ipython>=8.12.0"]

Hi Tirtho, could you add some more detail on how this fixed the issue for you? Did you uninstall the default jaxlib and then install jax[cpu] ?

from fbpinns.

I have the same issue. It also occurs when I run main.py in the multi-level paper.
Please advise

from fbpinns.

I faced the same issue. I fixed it by installing pip install -U "jax[cpu]" link, and rest of the dependencies "optax>=0.1.4", "numpy>=1.24.2", "scipy>=1.10.1", "matplotlib>=3.7.1", "tensorboardX>=2.6", "ipython>=8.12.0"]

Hi Tirtho, could you add some more detail on how this fixed the issue for you? Did you uninstall the default jaxlib and then install jax[cpu] ?

Hi, Yes, I have uninstalled everything and created a new python environment as suggested. Afterward I installed the jax[cpu] version and so on.

from fbpinns.

This looks like an issue with jax.ops.segment_sum and the internal jax operations not broadcasting correctly @jseeyave @avd28 @ruzhongliang @tirtho109 can you provide more details about your OS / python version / jax version / whether you are running on GPU or CPU? If we can produce a minimal version of this bug (a script which simply calls jax.ops.segment_sum with the same shape arrays which cause the error above) we can raise an issue in the main JAX repo

from fbpinns.

I'm running on CPU, OS is Ubuntu 20.04.5 LTS, Python 3.9.5, jax/jaxlib-0.4.25.
Please see screenshot below when I try to run '1. Defining your own problem - 1D harmonic oscillator'
'jax.ops.segment_sum' is called twice in trainers.py. I added 2 print statements after each call. They are in the screenshot between the [info] lines. e.g. "line 164. trainers.py. np.size(u)= 400"

Please let me know if there is anything else I can provide to help debug the issue.

from fbpinns.

I narrowed it down and found that there is inconsistency between the shapes of the output arrays of jnp.unique between different jax versions. I will upload a patch soon.

from fbpinns.

Ok I have patched, closing this issue as it should be fixed now.

from fbpinns.