iree-jax-gpt2-error.txt

(ajax) ✔ ~/w/iree-jax [main|✔]
19:53 $ python3 models/gpt2/export.py
I0423 19:54:00.877818 8498699072 xla_bridge.py:170] Remote TPU is not linked into jax; skipping remote TPU.
I0423 19:54:00.877909 8498699072 xla_bridge.py:353] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
I0423 19:54:00.877947 8498699072 xla_bridge.py:353] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0423 19:54:00.877974 8498699072 xla_bridge.py:353] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0423 19:54:00.878067 8498699072 xla_bridge.py:353] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
I0423 19:54:00.878116 8498699072 xla_bridge.py:353] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
Traceback (most recent call last):
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 115, in <module>
    absl.app.run(main)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 106, in main
    f.write(str(Program.get_mlir_module(module)))
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 377, in get_mlir_module
    info = Program.get_info(Program._get_instance(m))
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 372, in _get_instance
    m = m()
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 459, in __new__
    export_function()
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 454, in export_function
    info.export_module.def_func(invoke_with_self,
  File "/Users/y/w/iree-jax/iree/jax/exporter.py", line 206, in def_func
    return_py_value = f(*argument_py_tree)
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 452, in invoke_with_self
    return func_def.callable(self, *args, **kwargs)
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 61, in encode
    kv, x = self._encode(self._params, prompt, t)
  File "/Users/y/w/iree-jax/iree/jax/tracing.py", line 54, in __call__
    return current_ir_trace().handle_call(self, args, kwargs)
  File "/Users/y/w/iree-jax/iree/jax/tracing.py", line 114, in handle_call
    return target.resolve_call(self, *args, **kwargs)
  File "/Users/y/w/iree-jax/iree/jax/builtins.py", line 60, in resolve_call
    lowered = self.jit_f.lower(*abstract_args)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/api.py", line 743, in lower
    computation = dispatch.lower_xla_callable(
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/dispatch.py", line 472, in lower_xla_callable
    device = _xla_callable_device(nreps, backend, device, arg_devices)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/dispatch.py", line 651, in _xla_callable_device
    return xb.get_backend(backend).get_default_device_assignment(1)[0]
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/lib/xla_bridge.py", line 425, in get_backend
    return _get_backend_uncached(platform)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/jax/_src/lib/xla_bridge.py", line 417, in _get_backend_uncached
    raise RuntimeError(f"Unknown backend {platform}")
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: Unknown backend iree

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 115, in <module>
    absl.app.run(main)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/Users/y/miniforge3/envs/ajax/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 106, in main
    f.write(str(Program.get_mlir_module(module)))
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 377, in get_mlir_module
    info = Program.get_info(Program._get_instance(m))
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 372, in _get_instance
    m = m()
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 459, in __new__
    export_function()
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 454, in export_function
    info.export_module.def_func(invoke_with_self,
  File "/Users/y/w/iree-jax/iree/jax/exporter.py", line 206, in def_func
    return_py_value = f(*argument_py_tree)
  File "/Users/y/w/iree-jax/iree/jax/program_api.py", line 452, in invoke_with_self
    return func_def.callable(self, *args, **kwargs)
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 61, in encode
    kv, x = self._encode(self._params, prompt, t)
  File "/Users/y/w/iree-jax/iree/jax/tracing.py", line 54, in __call__
    return current_ir_trace().handle_call(self, args, kwargs)
  File "/Users/y/w/iree-jax/iree/jax/tracing.py", line 114, in handle_call
    return target.resolve_call(self, *args, **kwargs)
  File "/Users/y/w/iree-jax/iree/jax/builtins.py", line 60, in resolve_call
    lowered = self.jit_f.lower(*abstract_args)
RuntimeError: Unknown backend iree