Jax Version and Error

[Zhengwu-Miao]

Hello, I have the following questions:

1. What version of Jax are you using? When I ran PI_DeepONet_DR, I found that index_update and index were deprecated in newer versions of Jax, and optimizers were not included in lower versions of Jax.

"from jax.example_libraries import optimizers" "from jax.ops import index_update, index"

2. An error was found in PI_DeepONet_Burger.ipynb

The error occurs in the following code "vmap(generate_one_ics_training_data, in_axes=(0, 0, None, None))(keys, u0_train, m, P_ics_train) " The error reported is vmap got inconsistent sizes for array axes to be mapped:

• one axis had size 1000: axis 0 of argument key of type uint32[1000,2];
• one axis had size 10: axis 0 of argument u0 of type float32[10,101]

[rmojgani]

for 2.

I guess this is for the Burgers examples. There the number of input samples does not match with that of the MATLAB code. Either

N_train =10 #00      # number of input samples used for training

in PI_DeepONet_Burger or fix that in the matlab

[Zhengwu-Miao]

Thanks