[tab] slurm-13723.out (plaintext)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
Running on 16 devices.
2025-07-20 15:03:16.026999: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.036740: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.057628: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.069914: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.079912: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.079912: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.080022: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.080173: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.081358: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.080353: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.082585: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.084648: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.084677: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.089674: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.094441: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:16.095077: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:24.297069: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.658124: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.658122: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.706842: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.732398: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.732403: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.754489: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.772285: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.772289: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.785636: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.800901: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.814858: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.815606: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.822806: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:25.826878: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:31.164179: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:33.916472: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.612178: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.626415: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.674805: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.683095: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.692297: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.709654: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.744234: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.756363: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.770414: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.793038: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.869492: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.869492: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.881664: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:36.896683: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:37.218955: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
2025-07-20 15:03:39.969727: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:39.978146: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.040065: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.048869: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.049941: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.088329: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.103548: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.116496: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.145339: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.166537: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.229826: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.229827: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.241366: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:03:40.258171: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
2025-07-20 15:03:41.740483: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
2025-07-20 15:03:43.918230: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
Counting all components: ['tokenizer', 'lam', 'dynamics']
Parameter counts:
{'tokenizer': 33750256, 'lam': 17229792, 'dynamics': 26555904, 'total': 77535952}
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
WARNING:absl:Dropping 8 examples of 1000 examples (shard 16).
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `param-dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.float32'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/tyro/_parsers.py:347: UserWarning: The field `dtype` is annotated with type `<class 'numpy.dtype'>`, but the default value `<class 'jax.numpy.bfloat16'>` has type `<class 'jax._src.numpy.scalar_types._ScalarMeta'>`. We'll try to handle this gracefully, but it may cause unexpected behavior.
  warnings.warn(message)
2025-07-20 15:05:50.722915: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.723510: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.724343: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.725080: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.726687: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.726765: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.728527: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:50.989874: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.990457: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.991318: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.992083: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.993742: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.993798: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:50.995616: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:51.180777: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.181379: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.182211: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.183019: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.184791: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.184870: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.186697: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:51.319547: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.319909: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.320530: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.321081: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.322357: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.322380: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.323722: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.386548: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.387136: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.387974: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.388725: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.390350: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.390431: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.392169: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:51.496899: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
2025-07-20 15:05:51.497358: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.498150: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:51.498866: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.500449: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.500480: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:51.502179: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:52.265036: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.265669: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.266541: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.267334: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.268988: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.269088: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.270902: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:52.488050: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.488680: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.489522: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.490277: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.491841: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.491920: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:52.493736: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:57.564209: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.564899: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.565752: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.566529: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.568212: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.568295: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.570136: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:57.893641: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.894410: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.895318: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.896148: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.897874: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.897994: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.899845: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.908506: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.909034: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.909913: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.910703: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.912325: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.912426: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:57.914160: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:58.097966: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.098601: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.099494: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.100301: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.102037: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.102136: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.103962: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:58.454548: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.455251: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.456151: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.456954: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.458672: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.458768: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.460568: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.486614: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.487002: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.487617: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.488166: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.489485: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.489507: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.490866: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
2025-07-20 15:05:58.652824: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.653457: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.654309: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.655066: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.656722: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.656828: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.658583: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.720971: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.721363: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.721969: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.722564: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.723828: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.723845: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
2025-07-20 15:05:58.725218: W external/xla/xla/service/gpu/autotuning/dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. Maybe the hints set does not contain a good representative set of valid configs?Working around this by using the full hints set instead.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

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

Traceback (most recent call last):
  File "/fast/home/franz.srambical/jafar/train_dynamics.py", line 335, in <module>
    train_state, loss, recon, metrics = train_step(train_state, inputs)
jaxlib._jax.XlaRuntimeError: CANCELLED: Contracting dimension is too fragmented.
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 12 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '