Issue with using your code in Colab because of JAX! about learning-scientific_machine_learning_residual_based_attention_pinns_deeponets HOT 1 OPEN

Hi, thanks for reporting this issue.
I used the google colab GPU. However, it looks like they changed something in the new updates, which does not work now.
We only use "jax. experimental" to get our optimizer. I would recommend finding a way to use a previous version of Colab or importing the jax. experimental library. If not, as an alternative, you can try optax, but you may need to add some minor changes downstream too.
To install optax you would use the following commands:
!pip install optax
import optax
Then you can create an optimizer using the following lines:
optimizer = optax.adam(optax.exponential_decay(lr0, decay_step, decay_rate,)) opt_state = optimizer.init(params)
To update our parameters you need to compute the gradients of your loss. You can do something like this:
de los_fn(params):
'"your loss function here"
return loss
grad_fn = value_and_grad(loss_fn,)
loss, grads = grad_fn(params)
updates, opt_state = optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates)

Thanks for letting me know and I am sorry I could not help much this time.

from learning-scientific_machine_learning_residual_based_attention_pinns_deeponets.