Unexpected PJRT_Plugin_Attributes_Args size: expected 32, got 24.

Hello @alvarosg,

I've come across an error trying to run the `graphcast_demo.ipynb` notebook on a TPU `v5litepod-1`.

The following error is triggered when the rollout is performed, likely when the model is loaded or when the forward operator is called:

`pjrt_c_api_helpers.cc:258: Unexpected PJRT_Plugin_Attributes_Args size: expected 32, got 24. The plugin is likely built with a later version than the framework. This plugin is built with PJRT API version 0.75`

This error ultimately leads to an `Aborted (core dumped)`, stopping the rollout.

I feel that it has to do with my TPU not being compatible with the version of `jax` or `libtpu` libraries I have installed in my python venv.

Did you ever encounter this problem?

Thank you very much