search

PredictiveIntelligenceLab / Physics-informed-DeepONets

266 stars 83 forks source link

Unexpected arguments warning message about jax.vmap #4

Closed xuliang5115 closed 2 years ago

xuliang5115 commented 2 years ago

Hello Sifan, First of all, thank you for making this code public. I'm trying to run and understand the code of model PI_DeepONet for the purpose of calculating a parabolic equation problem. I rewritte the .ipyqb file into a model package to build the project structure, and then I noticed a warning message about vmap, but I didn't have any changes to the function operator_net.

Being cautious about the unknown, I would like to know whether this warning happend and the how to clean this warning. Many thanks! Physics-informed-DeepONets.zip

xuliang5115 commented 2 years ago

image

sifanexisted commented 2 years ago

Hi Liang,

We don't meet any warnings when running the notebook. It's better to provide more information of your warnings.

xuliang5115 commented 2 years ago

The warning only appears in Pycharm-2021.2.3 with jax~=0.2.24, there is no warnings when running the notebook. Sorry, I found the reason is not installed correctly, lack of jaxlab and then it seems that jupyterlab can replace it