Colab XLA Error in JAX_for_LLM_pretraining.ipynb

[Matt-Hurd]

On the LLM Pretraining Colab for step 11, I'm running into the following error.

XlaRuntimeError                           Traceback (most recent call last)
[<ipython-input-11-06e85c7e5c31>](https://localhost:8080/#) in <cell line: 0>()
      8 start_prompt = "Once upon a time"
      9 start_tokens = tokenizer.encode(start_prompt)[:maxlen]
---> 10 generated_text = model.generate_text(
     11     maxlen, start_tokens
     12 )

11 frames
    [... skipping hidden 11 frame]

[/usr/local/lib/python3.11/dist-packages/jax/_src/compiler.py](https://localhost:8080/#) in backend_compile(backend, module, options, host_callbacks)
    301     # TODO(sharadmv): remove this fallback when all backends allow `compile`
    302     # to take in `host_callbacks`
--> 303     return backend.compile(built_c, compile_options=options)
    304   except xc.XlaRuntimeError as e:
    305     for error_handler in _XLA_RUNTIME_ERROR_HANDLERS:

XlaRuntimeError: UNKNOWN: -:0:0: error: unknown vhlo type code: 33
-:0:0: note: in bytecode version 6 produced by: StableHLO_v1.7.1

This error started occurring sometime between Dec 1 and Jan 13.

Versions tested: 2025.1.9, 2024.12.10, 2024.11.1

[windmaple]

Seems specific to Colab TPU runtime. I did not see this error on Kaggle TPU or local CPU.

[windmaple]

Workaround is to install jaxtpu manually after JAX stack