Easydel support on TPU v4.8 - getting exception

[jchauhan]

Describe the bug

  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/EasyDel/modules/llama/llama_configuration.py", line 5, in <module>
    from ..easydel_modelling_utils import EasyDelPretrainedConfig
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/EasyDel/modules/easydel_modelling_utils.py", line 29, in <module>
    class EasyDelPretrainedConfig(PretrainedConfig):
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/EasyDel/modules/easydel_modelling_utils.py", line 79, in EasyDelPretrainedConfig
    backend: Optional[None] = jax.default_backend(),
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 894, in default_backend
    return get_backend(None).platform
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 823, in get_backend
    return _get_backend_uncached(platform)
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 803, in _get_backend_uncached
    bs = backends()
  File "/home/neo/research/belie/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 711, in backends
    raise RuntimeError(err_msg)
RuntimeError: Unable to initialize backend 'tpu': INTERNAL: Failed to get global TPU topology. (set JAX_PLATFORMS='' to automatically choose an available backend)

To Reproduce

1. Use tpu v4.8 VM
2. Install easydel and run an example such as llama

[erfanzar]

Error is not related to easydel that's from jax you have to update jax and tensorflow version to latest