Current thread 0x00007f8d2666e480 (most recent call first):
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py", line 619 in sharded_aval
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py", line 735 in lower_jaxpr_to_module
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 1832 in _cached_lowering_to_hlo
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2031 in lower_sharding_computation
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/profiler.py", line 340 in wrapper
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 188 in sharded_lowering
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 157 in xla_primitive_callable
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/util.py", line 256 in cached
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/util.py", line 263 in wrapper
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 128 in apply_primitive
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 868 in process_primitive
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 388 in bind_with_trace
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 385 in bind
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 555 in _convert_element_type
File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 2089 in array
File "/t9k/mnt/alphafold/alphafold/model/utils.py", line 116 in flat_params_to_haiku
File "/t9k/mnt/alphafold/alphafold/model/data.py", line 33 in get_model_haiku_params
File "/t9k/mnt/alphafold/./run_alphafold.py", line 518 in main
File "/t9k/mnt/.local/lib/python3.11/site-packages/absl/app.py", line 258 in _run_main
File "/t9k/mnt/.local/lib/python3.11/site-packages/absl/app.py", line 312 in run
File "/t9k/mnt/alphafold/./run_alphafold.py", line 570 in
Current thread 0x00007f8d2666e480 (most recent call first): File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py", line 619 in sharded_aval File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py", line 735 in lower_jaxpr_to_module File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 1832 in _cached_lowering_to_hlo File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2031 in lower_sharding_computation File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/profiler.py", line 340 in wrapper File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 188 in sharded_lowering File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 157 in xla_primitive_callable File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/util.py", line 256 in cached File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/util.py", line 263 in wrapper File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/dispatch.py", line 128 in apply_primitive File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 868 in process_primitive File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 388 in bind_with_trace File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/core.py", line 385 in bind File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 555 in _convert_element_type File "/t9k/mnt/.local/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 2089 in array File "/t9k/mnt/alphafold/alphafold/model/utils.py", line 116 in flat_params_to_haiku File "/t9k/mnt/alphafold/alphafold/model/data.py", line 33 in get_model_haiku_params File "/t9k/mnt/alphafold/./run_alphafold.py", line 518 in main File "/t9k/mnt/.local/lib/python3.11/site-packages/absl/app.py", line 258 in _run_main File "/t9k/mnt/.local/lib/python3.11/site-packages/absl/app.py", line 312 in run File "/t9k/mnt/alphafold/./run_alphafold.py", line 570 in