Graphcast error on Mac os

[seonyeongyang]

I'm trying to perform a graphcast model in a MacOS M1 environment, but I'm getting the following error:

To use the CPU, I set it up as follows.
pip install -r requirements.txt

This is the library version.
chex 0.1.85
jax 0.4.23
jaxlib 0.4.23

The same error occurs even if i run "pip install -upgrad jaxlib".

What causes the following error?

`2024-02-02 10:55:15,902 INFO Writing results to graphcast.grib.
2024-02-02 10:55:16,172 INFO Model description:
GraphCast model at 0.25deg resolution, with 13 pressure levels. This model is
trained on ERA5 data from 1979 to 2017, and fine-tuned on HRES-fc0 data from
2016 to 2021 and can be causally evaluated on 2022 and later years. This model
does not take total_precipitation_6hr as inputs and can make predictions in an
operational setting (i.e., initialised from HRES-fc0).

2024-02-02 10:55:16,172 INFO Model license:
The model weights are licensed under the Creative Commons
Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0). You
may obtain a copy of the License at:
https://creativecommons.org/licenses/by-nc-sa/4.0/.
The weights were trained on ERA5 data, see README for attribution statement.

2024-02-02 10:55:16,172 INFO Loading params/GraphCast_operational - ERA5-HRES 1979-2021 - resolution 0.25 - pressure levels 13 - mesh 2to6 - precipitation output only.npz: 0.3 second.
2024-02-02 10:55:16,172 INFO Building model: 0.3 second.
2024-02-02 10:55:16,172 INFO Loading surface fields from CDS
2024-02-02 10:55:16,301 INFO Loading pressure fields from CDS

2024-02-02 10:55:20,525 INFO Creating forcing variables: 4 seconds.
2024-02-02 10:55:22,446 INFO Converting GRIB to xarray: 1 second.
2024-02-02 10:55:27,549 INFO Reindexing: 5 seconds.
2024-02-02 10:55:27,665 INFO Creating training data: 11 seconds.
2024-02-02 10:55:37,317 INFO Extracting input targets: 9 seconds.
2024-02-02 10:55:37,317 INFO Creating input data (total): 21 seconds.
2024-02-02 10:55:37,349 INFO Unable to initialize backend 'cuda':
2024-02-02 10:55:37,350 INFO Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
2024-02-02 10:55:37,354 INFO Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/opt/homebrew/anaconda3/bin/../lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file), '/usr/local/lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache)
/Users/seonyeong/PycharmProjects/graphcast_wrf/.venv/lib/python3.11/site-packages/graphcast/autoregressive.py:202: FutureWarning: The return type of Dataset.dims will be changed to return a set of dimension names in future, in order to be more consistent with DataArray.dims. To access a mapping from dimension names to lengths, please use Dataset.sizes.
scan_length = targets_template.dims['time']
/Users/seonyeong/PycharmProjects/graphcast_wrf/.venv/lib/python3.11/site-packages/graphcast/autoregressive.py:115: FutureWarning: The return type of Dataset.dims will be changed to return a set of dimension names in future, in order to be more consistent with DataArray.dims. To access a mapping from dimension names to lengths, please use Dataset.sizes.
num_inputs = inputs.dims['time']
zsh: killed ai-models --input cds --date 20230601 --time 0000 --lead-time 120 graphcast`

[Admar17]

The first thing I noticed is that you are trying to call in a TPU backend, which isn't available on Apple Silicon (M1,M2, and M3). Have you tried updating your code to only use the CPU instance on your Mac?