ASAN + poisoning for https://github.com/openxla/openxla-pjrt-plugin/issues/170

paxmlsegfault_asan_poison.txt

2023-06-05 17:37:07.425028: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.664716: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.664797: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.664863: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.664927: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.696031: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu
2023-06-05 17:37:15.696218: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
[IREE-PJRT] DEBUG: Using IREE compiler binary: /usr/local/lib/python3.10/dist-packages/iree/compiler/_mlir_libs/libIREECompiler.so
[IREE-PJRT] DEBUG: Compiler Version: 20230604.542 @ 7400c8546f33cc76680c5e235d17617a5ec8c18c (API version 1.2)
[IREE-PJRT] DEBUG: Using partitioner binary: /workspace/openxla-pjrt-plugin/bazel-bin/partitioner/libOpenXLAPartitioner.so
[IREE-PJRT] DEBUG: Partitioner version: <unknown> (API version 1.1)
[IREE-PJRT] DEBUG: CUDA driver created
I0605 17:37:15.740515 139756486557120 setup_jax.py:72] JAX process: 0 / 1
I0605 17:37:15.740673 139756486557120 setup_jax.py:73] JAX devices: [GPU-b0fbccec-7593-9c0c-35de-cbfc04b9d09a]
I0605 17:37:15.740910 139756486557120 setup_jax.py:74] jax.device_count(): 1
I0605 17:37:15.741026 139756486557120 setup_jax.py:75] jax.local_device_count(): 1
I0605 17:37:15.741060 139756486557120 setup_jax.py:76] jax.process_count(): 1
Registered experiment `paxml.tasks.lm.params.lm_cloud.LargeMlp`
Registered experiment `paxml.tasks.lm.params.lm_cloud.SmallMlp`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudTransformerAdam`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudTransformerAdamTest`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudTransformerAdamLimitSteps`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdTest`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd2B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd32B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd64B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd128B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd256B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd512B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmd1024B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdPipeline9B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdPipeline175B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdMultislice2B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdPipelineMultislice2B`
Registered experiment `paxml.tasks.lm.params.lm_cloud.LmCloudSpmdPipelineMultislice2BCircular`
Registered experiment `paxml.tasks.lm.params.c4.LmCloudSpmdAdam`
Registered experiment `paxml.tasks.lm.params.c4.LmCloudSpmdAdamLimitSteps`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdAdam`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdGpt3AdamOrgHPBS1p5k1536Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineAdam`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamOrgHPBS1p5k768Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamMLPerfHPBS1p5k768Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamMLPerfHPBS2k512Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamMLPerfHPBS3k768Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamMLPerfHPBS4k1024Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3AdamMLPerfHPBS8k1024Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4Spmd1BAdam4Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4Spmd1BAdam4ReplicasLimitSteps`
Registered experiment `paxml.tasks.lm.params.c4.C4Spmd2BAdam4Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4Spmd16BAdam32Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4Spmd32BAdam64Replicas`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdGpt3L16AdamOrgHP`
Registered experiment `paxml.tasks.lm.params.c4.C4SpmdPipelineGpt3SmallAdam8Replicas`
W0605 17:37:16.975340 139756486557120 gpu_fast_attention.py:41] jax_triton not found, please `pip install jax-triton`
Registered experiment `tasks.lm.params.nvidia.NVIDIA1_3B`
Registered experiment `tasks.lm.params.nvidia.NVIDIA5B`
Registered experiment `tasks.lm.params.nvidia.NVIDIA8_3B`
Registered experiment `tasks.lm.params.nvidia.NVIDIA10B`
Registered experiment `tasks.lm.params.nvidia.NVIDIA40BProxy`
Registered experiment `tasks.lm.params.nvidia.NVIDIA70BProxy`
Registered experiment `tasks.lm.params.nvidia.NVIDIA116BProxy`
Registered experiment `tasks.lm.params.nvidia.NVIDIA175BProxy`
Registered experiment `tasks.lm.params.nvidia.TestSmallConfig`
Registered experiment `tasks.lm.params.nvidia.NVIDIA1_3BPmap`
I0605 17:37:16.984491 139756486557120 local.py:45] Setting task status: process_index: 0, process_count: 1
I0605 17:37:16.984750 139756486557120 local.py:50] Created artifact job_log_dir of type ArtifactType.DIRECTORY and value log_NVIDIA1_3BPmap.
I0605 17:37:17.464085 139756486557120 local.py:45] Setting task status: Train experiment tasks.lm.params.nvidia.NVIDIA1_3BPmap at log_NVIDIA1_3BPmap
I0605 17:37:17.464244 139756486557120 train.py:139] [PAX STATUS] Starting `train_and_evaluate`
I0605 17:37:17.601573 139756486557120 train.py:146] [PAX STATUS] Obtaining and initializing datasets.
I0605 17:37:17.604288 139756486557120 train.py:162] [PAX STATUS]: Done initializing dataset objects
I0605 17:37:17.604343 139756486557120 train.py:164] train_input_p:
I0605 17:37:17.604944 139756486557120 train.py:168]   allow_fixed_file_random_seed : False
I0605 17:37:17.604993 139756486557120 train.py:168]   batch_padding_size : 0
I0605 17:37:17.605027 139756486557120 train.py:168]   batch_size : NoneType
I0605 17:37:17.605060 139756486557120 train.py:168]   cls : type/praxis.base_input/LingvoInputAdaptor
I0605 17:37:17.605092 139756486557120 train.py:168]   cluster_do_eval : False
I0605 17:37:17.605139 139756486557120 train.py:168]   custom_device_order : NoneType
I0605 17:37:17.605173 139756486557120 train.py:168]   eval_loop_num_batches : 1
I0605 17:37:17.605206 139756486557120 train.py:168]   experimental_remote_input : False
I0605 17:37:17.605238 139756486557120 train.py:168]   infeed_host_index : 0
I0605 17:37:17.605269 139756486557120 train.py:168]   input.activation_split_dims_mapping : NoneType
I0605 17:37:17.605301 139756486557120 train.py:168]   input.add_name_to_theta : False
I0605 17:37:17.605332 139756486557120 train.py:168]   input.allow_implicit_capture : NoneType
I0605 17:37:17.605364 139756486557120 train.py:168]   input.batch_size : 1
I0605 17:37:17.605396 139756486557120 train.py:168]   input.cls : type/paxml.tasks.lm.input_generator/SyntheticLmData
I0605 17:37:17.605428 139756486557120 train.py:168]   input.decoder_samples_per_summary : NoneType
I0605 17:37:17.605461 139756486557120 train.py:168]   input.device_mesh : NoneType
I0605 17:37:17.605493 139756486557120 train.py:168]   input.dtype : float32
I0605 17:37:17.605524 139756486557120 train.py:168]   input.eval_samples_per_summary : NoneType
I0605 17:37:17.605556 139756486557120 train.py:168]   input.file_datasource : NoneType
I0605 17:37:17.605586 139756486557120 train.py:168]   input.filter_sparse_tensors : False
I0605 17:37:17.605615 139756486557120 train.py:168]   input.fprop_dtype : NoneType
I0605 17:37:17.605644 139756486557120 train.py:168]   input.inference_driver_name : NoneType
I0605 17:37:17.605674 139756486557120 train.py:168]   input.input_stats_summary_interval_steps : 10
I0605 17:37:17.605703 139756486557120 train.py:168]   input.is_inference : NoneType
I0605 17:37:17.605733 139756486557120 train.py:168]   input.name : 'input'
I0605 17:37:17.605762 139756486557120 train.py:168]   input.num_partitions : NoneType
I0605 17:37:17.605792 139756486557120 train.py:168]   input.num_samples : 0
I0605 17:37:17.605822 139756486557120 train.py:168]   input.outfeed_in_logical_order : False
I0605 17:37:17.605851 139756486557120 train.py:168]   input.params_init.custom_v_init : NoneType
I0605 17:37:17.605880 139756486557120 train.py:168]   input.params_init.method : 'xavier'
I0605 17:37:17.605910 139756486557120 train.py:168]   input.params_init.scale : 1.000001
I0605 17:37:17.605938 139756486557120 train.py:168]   input.params_init.seed : NoneType
I0605 17:37:17.605968 139756486557120 train.py:168]   input.random_seed : NoneType
I0605 17:37:17.605997 139756486557120 train.py:168]   input.remote.max_inflights_per_target : 32
I0605 17:37:17.606026 139756486557120 train.py:168]   input.resettable : False
I0605 17:37:17.606055 139756486557120 train.py:168]   input.seq_len : 2048
I0605 17:37:17.606085 139756486557120 train.py:168]   input.skip_lp_regularization : NoneType
I0605 17:37:17.606114 139756486557120 train.py:168]   input.tpu_embedding_mode : 'train'
I0605 17:37:17.606144 139756486557120 train.py:168]   input.tpu_infeed_parallelism : 1
I0605 17:37:17.606173 139756486557120 train.py:168]   input.use_partitioned_infeed_queue : False
I0605 17:37:17.606202 139756486557120 train.py:168]   input.use_per_core_infeed : False
I0605 17:37:17.606231 139756486557120 train.py:168]   input.use_per_host_infeed : False
I0605 17:37:17.606261 139756486557120 train.py:168]   input.vn.deterministic : NoneType
I0605 17:37:17.606290 139756486557120 train.py:168]   input.vn.global_vn : False
I0605 17:37:17.606320 139756486557120 train.py:168]   input.vn.per_step_vn : False
I0605 17:37:17.606349 139756486557120 train.py:168]   input.vn.scale : NoneType
I0605 17:37:17.606379 139756486557120 train.py:168]   input.vn.seed : NoneType
I0605 17:37:17.606409 139756486557120 train.py:168]   input.vn.start_step : 0
I0605 17:37:17.606438 139756486557120 train.py:168]   input.weight_split_dims_mapping : NoneType
I0605 17:37:17.606467 139756486557120 train.py:168]   input_checkpointing_enabled : False
I0605 17:37:17.606497 139756486557120 train.py:168]   input_random_seed : NoneType
I0605 17:37:17.606526 139756486557120 train.py:168]   is_training : True
I0605 17:37:17.606556 139756486557120 train.py:168]   name : ''
I0605 17:37:17.606590 139756486557120 train.py:168]   num_batches : NoneType
I0605 17:37:17.606620 139756486557120 train.py:168]   num_infeed_hosts : 0
I0605 17:37:17.606650 139756486557120 train.py:168]   reset_for_eval : False
I0605 17:37:17.606679 139756486557120 train.py:168]   tf_data_service_address : NoneType
I0605 17:37:17.606710 139756486557120 train.py:169] task_p:
I0605 17:37:17.628924 139756486557120 train.py:171]   cls : type/paxml.tasks_lib/SingleTask
I0605 17:37:17.629092 139756486557120 train.py:171]   decode.cls : type/paxml.tasks_lib/SingleTask.Decode
I0605 17:37:17.629150 139756486557120 train.py:171]   decode.prng_key_fold_with_batch_index : False
I0605 17:37:17.629184 139756486557120 train.py:171]   decode.prng_key_fold_with_global_step : True
I0605 17:37:17.629215 139756486557120 train.py:171]   decode.random_seed : 1234
I0605 17:37:17.629246 139756486557120 train.py:171]   early_stopping_fn : NoneType
I0605 17:37:17.629276 139756486557120 train.py:171]   evaluate.apply_mutable_list : ['aux_loss', 'summaries', 'non_trainable']
I0605 17:37:17.629308 139756486557120 train.py:171]   evaluate.cls : type/paxml.tasks_lib/SingleTask.Evaluate
I0605 17:37:17.629338 139756486557120 train.py:171]   evaluate.random_seed : 1234
I0605 17:37:17.629369 139756486557120 train.py:171]   infer.cls : type/paxml.tasks_lib/SingleTask.Infer
I0605 17:37:17.629403 139756486557120 train.py:171]   infer.random_seed : 1234
I0605 17:37:17.629433 139756486557120 train.py:171]   infer_writer : NoneType
I0605 17:37:17.629464 139756486557120 train.py:171]   loss_aggregator : NoneType
I0605 17:37:17.629495 139756486557120 train.py:171]   metrics : NoneType
I0605 17:37:17.629525 139756486557120 train.py:171]   model.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.629556 139756486557120 train.py:171]   model.apply_eval_sample_weights : False
I0605 17:37:17.629586 139756486557120 train.py:171]   model.cls : type/praxis.layers.models/LanguageModel
I0605 17:37:17.629616 139756486557120 train.py:171]   model.contiguous_submeshes : NoneType
I0605 17:37:17.629646 139756486557120 train.py:171]   model.count_tokens : False
I0605 17:37:17.629676 139756486557120 train.py:171]   model.dcn_mesh_shape : NoneType
I0605 17:37:17.629707 139756486557120 train.py:171]   model.decoder_tpl.cls : type/praxis.decoder_hparams/GreedyDecoderHParams
I0605 17:37:17.629737 139756486557120 train.py:171]   model.decoder_tpl.decode_loop_mesh_axes_transpose : NoneType
I0605 17:37:17.629767 139756486557120 train.py:171]   model.decoder_tpl.emb_lookup_style : 'matmul'
I0605 17:37:17.629796 139756486557120 train.py:171]   model.decoder_tpl.eos_id : 2
I0605 17:37:17.629826 139756486557120 train.py:171]   model.decoder_tpl.fprop_for_prefix : False
I0605 17:37:17.629856 139756486557120 train.py:171]   model.decoder_tpl.lazy_prefix_broadcast : False
I0605 17:37:17.629886 139756486557120 train.py:171]   model.decoder_tpl.max_decode_steps : NoneType
I0605 17:37:17.629916 139756486557120 train.py:171]   model.decoder_tpl.min_prefix_len : 5
I0605 17:37:17.629946 139756486557120 train.py:171]   model.decoder_tpl.process_result_fn : NoneType
I0605 17:37:17.629976 139756486557120 train.py:171]   model.decoder_tpl.seqlen : 0
I0605 17:37:17.630006 139756486557120 train.py:171]   model.dtype : type/jax.numpy/float32
I0605 17:37:17.630036 139756486557120 train.py:171]   model.fprop_dtype : dtype[float32]
I0605 17:37:17.630066 139756486557120 train.py:171]   model.ici_mesh_shape : NoneType
I0605 17:37:17.630096 139756486557120 train.py:171]   model.lm_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.630126 139756486557120 train.py:171]   model.lm_tpl.cls : type/praxis.layers.transformer_models/TransformerLm
I0605 17:37:17.630156 139756486557120 train.py:171]   model.lm_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.630186 139756486557120 train.py:171]   model.lm_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.630216 139756486557120 train.py:171]   model.lm_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.630246 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.630275 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.cls : type/praxis.layers.normalizations/LayerNorm
I0605 17:37:17.630306 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.630335 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.630365 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.dim : 0
I0605 17:37:17.630395 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.630425 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.epsilon : 1e-06
I0605 17:37:17.630455 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.fprop_dtype : NoneType
I0605 17:37:17.630485 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.630515 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.mesh_axis_names : NoneType
I0605 17:37:17.630545 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.name : NoneType
I0605 17:37:17.630575 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.630605 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.params_init.method : 'xavier'
I0605 17:37:17.630635 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.params_init.scale : 1.000001
I0605 17:37:17.630665 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.reductions_in_fp32 : False
I0605 17:37:17.630694 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.630724 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.630754 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.use_bias : True
I0605 17:37:17.630784 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.use_scale : True
I0605 17:37:17.630813 139756486557120 train.py:171]   model.lm_tpl.final_ln_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.630843 139756486557120 train.py:171]   model.lm_tpl.fprop_dtype : NoneType
I0605 17:37:17.630873 139756486557120 train.py:171]   model.lm_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.630903 139756486557120 train.py:171]   model.lm_tpl.mesh_axis_names : NoneType
I0605 17:37:17.630933 139756486557120 train.py:171]   model.lm_tpl.model_dims : 2048
I0605 17:37:17.630962 139756486557120 train.py:171]   model.lm_tpl.model_type : 'causal'
I0605 17:37:17.630992 139756486557120 train.py:171]   model.lm_tpl.name : NoneType
I0605 17:37:17.631022 139756486557120 train.py:171]   model.lm_tpl.ngrammer_tpl : NoneType
I0605 17:37:17.631052 139756486557120 train.py:171]   model.lm_tpl.packed_input : True
I0605 17:37:17.631083 139756486557120 train.py:171]   model.lm_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.631113 139756486557120 train.py:171]   model.lm_tpl.params_init.method : 'xavier'
I0605 17:37:17.631142 139756486557120 train.py:171]   model.lm_tpl.params_init.scale : 1.000001
I0605 17:37:17.631172 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.activation_split_dims_mapping.emb_out_split_dims_mapping : NoneType
I0605 17:37:17.631202 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.631233 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.631264 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.cls : type/praxis.layers.base_ops/ArrayLookup
I0605 17:37:17.631293 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.631323 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.631354 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.631383 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.fprop_dtype : NoneType
I0605 17:37:17.631413 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.631443 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.mesh_axis_names : NoneType
I0605 17:37:17.631474 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.name : NoneType
I0605 17:37:17.631504 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.631534 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.params_init.method : 'xavier'
I0605 17:37:17.631564 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.params_init.scale : 1.000001
I0605 17:37:17.631594 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.631624 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.631654 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.array_lookup_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.631683 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.cls : type/praxis.layers.embedding_softmax/TrainablePositionalEmbedding
I0605 17:37:17.631715 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.631745 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.631775 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.631805 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.631835 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.631865 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.631894 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.631924 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.631954 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.631984 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.632014 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.632044 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.632073 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.632104 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.632134 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.632164 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.632194 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.632224 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.632253 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.embedding_dims : 0
I0605 17:37:17.632283 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.fprop_dtype : NoneType
I0605 17:37:17.632313 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.632343 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.lookup_style : 'matmul'
I0605 17:37:17.632373 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.max_seq_length : 2048
I0605 17:37:17.632403 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.max_timescale : 10000
I0605 17:37:17.632433 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.mesh_axis_names : NoneType
I0605 17:37:17.632463 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.min_timescale : 1
I0605 17:37:17.632493 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.name : NoneType
I0605 17:37:17.632523 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.632553 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.params_init.method : 'xavier'
I0605 17:37:17.632583 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.params_init.scale : 1.000001
I0605 17:37:17.632613 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.632643 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.632673 139756486557120 train.py:171]   model.lm_tpl.position_emb_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.632703 139756486557120 train.py:171]   model.lm_tpl.post_attention_ngrammer_tpls : NoneType
I0605 17:37:17.632733 139756486557120 train.py:171]   model.lm_tpl.record_activations_in_xent_output : False
I0605 17:37:17.632762 139756486557120 train.py:171]   model.lm_tpl.separate_embedding_tpl : NoneType
I0605 17:37:17.632792 139756486557120 train.py:171]   model.lm_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.632822 139756486557120 train.py:171]   model.lm_tpl.skip_aux_loss : False
I0605 17:37:17.632852 139756486557120 train.py:171]   model.lm_tpl.skip_compute_loss : False
I0605 17:37:17.632882 139756486557120 train.py:171]   model.lm_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.632912 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.activation_split_dims_mapping.emb_out_split_dims_mapping : NoneType
I0605 17:37:17.632942 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.632971 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.633002 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.cls : type/praxis.layers.base_ops/ArrayLookup
I0605 17:37:17.633031 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.633061 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.633091 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.633129 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.fprop_dtype : NoneType
I0605 17:37:17.633161 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.633191 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.mesh_axis_names : NoneType
I0605 17:37:17.633221 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.name : NoneType
I0605 17:37:17.633251 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.633281 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.params_init.method : 'xavier'
I0605 17:37:17.633311 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.params_init.scale : 1.000001
I0605 17:37:17.633340 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.633370 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.633401 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.array_lookup_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.633431 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.bi_tempered_loss_tpl : NoneType
I0605 17:37:17.633461 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.bias_init : 0.0
I0605 17:37:17.633491 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.cls : type/praxis.layers.embedding_softmax/SharedEmbeddingSoftmax
I0605 17:37:17.633521 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.633551 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.633581 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.633611 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.633641 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.633671 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.633702 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.633732 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.633762 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.633792 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.633822 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.633852 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.633882 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.633912 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.633942 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.633972 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.634001 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.634031 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.634061 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.634091 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.634122 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.cls : type/praxis.layers.activations/ReLU
I0605 17:37:17.634152 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.634182 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.634212 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.634242 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.fprop_dtype : NoneType
I0605 17:37:17.634272 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.634302 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.mesh_axis_names : NoneType
I0605 17:37:17.634332 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.name : NoneType
I0605 17:37:17.634362 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.634393 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.params_init.method : 'xavier'
I0605 17:37:17.634423 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.params_init.scale : 1.000001
I0605 17:37:17.634453 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.634483 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.634512 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.activation_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.634542 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.bias_init : 0.0
I0605 17:37:17.634572 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.cls : type/praxis.layers.linears/FeedForward
I0605 17:37:17.634603 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.634633 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.634662 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.634692 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.fprop_dtype : NoneType
I0605 17:37:17.634722 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.has_bias : True
I0605 17:37:17.634752 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.634781 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.input_dims : 0
I0605 17:37:17.634811 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.634842 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.cls : type/praxis.layers.linears/Linear
I0605 17:37:17.634871 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.634902 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.634932 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.634961 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.634991 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.635022 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.635052 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.635082 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.635112 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.635141 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.635172 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.635202 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.635232 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.635263 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.635293 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.635324 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.635354 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.635384 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.635414 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.fprop_dtype : NoneType
I0605 17:37:17.635444 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.635474 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.input_dims : 0
I0605 17:37:17.635504 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.mesh_axis_names : NoneType
I0605 17:37:17.635534 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.name : NoneType
I0605 17:37:17.635565 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.output_dims : 0
I0605 17:37:17.635594 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.635624 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.params_init.method : 'xavier'
I0605 17:37:17.635654 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.params_init.scale : 1.000001
I0605 17:37:17.635684 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.635714 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.635744 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.weight_init : NoneType
I0605 17:37:17.635774 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.linear_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.635804 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.mesh_axis_names : NoneType
I0605 17:37:17.635834 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.name : NoneType
I0605 17:37:17.635864 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.output_dims : 0
I0605 17:37:17.635895 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.635925 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.params_init.method : 'xavier'
I0605 17:37:17.635955 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.params_init.scale : 1.000001
I0605 17:37:17.635985 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.636015 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.636045 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.weight_init : NoneType
I0605 17:37:17.636075 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.feed_forward_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.636105 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.fprop_dtype : NoneType
I0605 17:37:17.636136 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.636165 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.input_dims : 0
I0605 17:37:17.636195 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.label_smoothing_apply_for_eval : True
I0605 17:37:17.636226 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.label_smoothing_prob : 0.0
I0605 17:37:17.636256 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.lookup_style : 'index'
I0605 17:37:17.636286 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.mesh_axis_names : NoneType
I0605 17:37:17.636316 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.name : NoneType
I0605 17:37:17.636346 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.num_classes : 0
I0605 17:37:17.636376 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.params_init.method : 'gaussian'
I0605 17:37:17.636406 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.params_init.scale : 0.022097086912079608
I0605 17:37:17.636436 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.scale_sqrt_depth : True
I0605 17:37:17.636466 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.636496 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.636526 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.soft_cap_logits : 30.0
I0605 17:37:17.636556 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.636586 139756486557120 train.py:171]   model.lm_tpl.softmax_tpl.z_loss_weight : 0.0
I0605 17:37:17.636616 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.636646 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.636676 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.atten_dropout_prob : NoneType
I0605 17:37:17.636706 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.cls : type/praxis.layers.transformers/StackedTransformer
I0605 17:37:17.636736 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.contiguous_submeshes : NoneType
I0605 17:37:17.636766 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.dcn_mesh_shape : NoneType
I0605 17:37:17.636796 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.dim_per_head : 64
I0605 17:37:17.636826 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.dropout_prob : 0.0
I0605 17:37:17.636857 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.dtype : type/jax.numpy/float32
I0605 17:37:17.636888 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.fold_padding_with_segment_mask : False
I0605 17:37:17.636918 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.fprop_dtype : NoneType
I0605 17:37:17.636949 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.gating_func : 'top2'
I0605 17:37:17.636978 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.hidden_dims : 8192
I0605 17:37:17.637008 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.ici_mesh_shape : NoneType
I0605 17:37:17.637038 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.input_dropout_prob : 0.0
I0605 17:37:17.637068 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.mask_self_attention : False
I0605 17:37:17.637098 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.mesh_axis_names : NoneType
I0605 17:37:17.637149 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.min_group_size : NoneType
I0605 17:37:17.637181 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.model_dims : 2048
I0605 17:37:17.637211 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.egch : NoneType
I0605 17:37:17.637242 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.egcm : NoneType
I0605 17:37:17.637273 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gec : NoneType
I0605 17:37:17.637304 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gecm : NoneType
I0605 17:37:17.637334 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gecs : NoneType
I0605 17:37:17.637365 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gs : NoneType
I0605 17:37:17.637395 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gsec : NoneType
I0605 17:37:17.637425 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.gsm : NoneType
I0605 17:37:17.637456 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.637486 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.637517 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.cls : type/praxis.layers.activations/ReLU
I0605 17:37:17.637548 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.637578 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.637609 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.637638 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.fprop_dtype : NoneType
I0605 17:37:17.637669 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.637699 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.mesh_axis_names : NoneType
I0605 17:37:17.637729 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.name : NoneType
I0605 17:37:17.637759 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.637790 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.params_init.method : 'xavier'
I0605 17:37:17.637820 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.params_init.scale : 1.000001
I0605 17:37:17.637851 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.637882 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.637912 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.activation_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.637943 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.add_skip_connection : True
I0605 17:37:17.637974 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.apply_padding_first : False
I0605 17:37:17.638004 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.cls : type/praxis.layers.transformers/TransformerFeedForwardMoe
I0605 17:37:17.638034 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.638064 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.638095 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.638125 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.expert_capacity_dim : 0
I0605 17:37:17.638155 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.expert_weight_shards : 1
I0605 17:37:17.638185 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.explicit_fan_in_fan_out_axes : False
I0605 17:37:17.638216 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.fprop_dtype : NoneType
I0605 17:37:17.638245 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.gating_func : 'top2'
I0605 17:37:17.638275 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.gating_logit_cap : 0.0
I0605 17:37:17.638305 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.hidden_dims : 0
I0605 17:37:17.638335 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.638365 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.input_dims : 0
I0605 17:37:17.638394 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.internal_gshard_variance_scaling_fan_in_init : True
I0605 17:37:17.638424 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.638454 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.cls : type/praxis.layers.normalizations/LayerNorm
I0605 17:37:17.638485 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.638515 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.638545 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.dim : 0
I0605 17:37:17.638575 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.638604 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.epsilon : 1e-06
I0605 17:37:17.638634 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.fprop_dtype : NoneType
I0605 17:37:17.638664 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.638695 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.mesh_axis_names : NoneType
I0605 17:37:17.638724 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.name : NoneType
I0605 17:37:17.638754 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.638784 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.params_init.method : 'xavier'
I0605 17:37:17.638813 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.params_init.scale : 1.000001
I0605 17:37:17.638843 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.reductions_in_fp32 : False
I0605 17:37:17.638873 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.638903 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.638933 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.use_bias : True
I0605 17:37:17.638963 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.use_scale : True
I0605 17:37:17.638993 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.ln_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.639023 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.mesh_axis_names : NoneType
I0605 17:37:17.639053 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.min_group_size : NoneType
I0605 17:37:17.639083 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.moe_gating_embedding_level : 'token'
I0605 17:37:17.639112 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.moe_load_balance_loss_weight : 1.0
I0605 17:37:17.639142 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.name : NoneType
I0605 17:37:17.639172 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.norm_policy : 'pre'
I0605 17:37:17.639203 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.num_experts : 0
I0605 17:37:17.639233 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.num_groups : 0
I0605 17:37:17.639262 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.639292 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.params_init.method : 'xavier'
I0605 17:37:17.639322 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.params_init.scale : 1.000001
I0605 17:37:17.639352 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_prob : 0.0
I0605 17:37:17.639382 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.639413 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.639443 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.639473 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.639504 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.dropout_at_eval : False
I0605 17:37:17.639533 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.639564 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.639594 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.639624 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.keep_prob : 1.0
I0605 17:37:17.639654 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.639684 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.name : NoneType
I0605 17:37:17.639714 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.noise_shape : NoneType
I0605 17:37:17.639744 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.639774 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.639804 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.639834 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.639864 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.639894 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.639924 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.transpose_qk : False
I0605 17:37:17.639954 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.relu_dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.639984 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_prob : 0.0
I0605 17:37:17.640014 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.640045 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.640075 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.640106 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.640136 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.dropout_at_eval : False
I0605 17:37:17.640166 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.640196 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.640226 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.640257 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.keep_prob : 1.0
I0605 17:37:17.640287 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.640317 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.name : NoneType
I0605 17:37:17.640347 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.noise_shape : NoneType
I0605 17:37:17.640377 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.640407 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.640438 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.640468 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.640498 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.640528 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.640558 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.transpose_qk : False
I0605 17:37:17.640588 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.640617 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_droppath_prob : 0.0
I0605 17:37:17.640647 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.residual_weight : 1.0
I0605 17:37:17.640677 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.second_expert_policy : 'all'
I0605 17:37:17.640707 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.640737 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.640768 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.unadjusted_expert_capacity_factor : 2.0
I0605 17:37:17.640797 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.use_gated_activation : False
I0605 17:37:17.640828 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.weight_split_dims_mapping.ehm : NoneType
I0605 17:37:17.640858 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.weight_split_dims_mapping.emh : NoneType
I0605 17:37:17.640888 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.weight_split_dims_mapping.me : NoneType
I0605 17:37:17.640918 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.moe_layer_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.640948 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.name : NoneType
I0605 17:37:17.640979 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.ngrammer_tpls : NoneType
I0605 17:37:17.641009 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.num_experts : 0
I0605 17:37:17.641039 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.num_groups : 1
I0605 17:37:17.641069 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.num_heads : 32
I0605 17:37:17.641099 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.num_layers : 1
I0605 17:37:17.641137 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.packed_input : False
I0605 17:37:17.641168 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.641199 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.params_init.method : 'xavier'
I0605 17:37:17.641229 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.params_init.scale : 1.000001
I0605 17:37:17.641259 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.relu_dropout_prob : NoneType
I0605 17:37:17.641289 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.residual_dropout_prob : NoneType
I0605 17:37:17.641318 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.residual_droppath_prob : 0.0
I0605 17:37:17.641348 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.shared_weight_layer_id : NoneType
I0605 17:37:17.641378 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.skip_lp_regularization : NoneType
I0605 17:37:17.641408 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.641438 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.allow_skip_cross_attention : False
I0605 17:37:17.641468 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.atten_dropout_prob : 0.0
I0605 17:37:17.641499 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.cls : type/praxis.layers.transformers/Transformer
I0605 17:37:17.641529 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.641559 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.cross_atten_tpl : NoneType
I0605 17:37:17.641589 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.641619 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dim_per_head : NoneType
I0605 17:37:17.641650 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.641684 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.641715 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.641745 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.641775 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.dropout_at_eval : False
I0605 17:37:17.641805 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.641834 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.641865 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.641894 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.keep_prob : 1.0
I0605 17:37:17.641924 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.641954 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.name : NoneType
I0605 17:37:17.641984 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.noise_shape : NoneType
I0605 17:37:17.642015 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.642046 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.642076 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.642107 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.642137 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.642167 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.642197 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.transpose_qk : False
I0605 17:37:17.642227 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.642257 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.642287 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.fprop_dtype : NoneType
I0605 17:37:17.642318 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.hidden_dims : 0
I0605 17:37:17.642349 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.642379 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.input_dims : 0
I0605 17:37:17.642410 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.642441 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.cls : type/praxis.layers.normalizations/LayerNorm
I0605 17:37:17.642472 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.642502 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.642532 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.dim : 0
I0605 17:37:17.642563 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.642593 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.epsilon : 1e-06
I0605 17:37:17.642624 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.fprop_dtype : NoneType
I0605 17:37:17.642654 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.642684 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.mesh_axis_names : NoneType
I0605 17:37:17.642715 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.name : NoneType
I0605 17:37:17.642745 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.642776 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.params_init.method : 'xavier'
I0605 17:37:17.642806 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.params_init.scale : 1.000001
I0605 17:37:17.642836 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.reductions_in_fp32 : False
I0605 17:37:17.642866 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.642896 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.642927 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.use_bias : True
I0605 17:37:17.642957 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.use_scale : True
I0605 17:37:17.642987 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ln_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.643017 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.mesh_axis_names : NoneType
I0605 17:37:17.643048 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.name : NoneType
I0605 17:37:17.643078 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.ngrammer_tpl : NoneType
I0605 17:37:17.643109 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.norm_policy : 'pre'
I0605 17:37:17.643139 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.num_heads : NoneType
I0605 17:37:17.643170 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.packed_input : False
I0605 17:37:17.643200 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.643230 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.params_init.method : 'xavier'
I0605 17:37:17.643261 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.params_init.scale : 1.000001
I0605 17:37:17.643291 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.relu_dropout_prob : 0.0
I0605 17:37:17.643321 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.residual_dropout_prob : 0.0
I0605 17:37:17.643351 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.residual_droppath_prob : 0.0
I0605 17:37:17.643382 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.643412 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.643442 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.activation_split_dims_mapping.bld : NoneType
I0605 17:37:17.643473 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.activation_split_dims_mapping.blnh : NoneType
I0605 17:37:17.643503 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.643534 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.atten_dropout_prob : 0.0
I0605 17:37:17.643564 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.atten_logit_cap : 50.0
I0605 17:37:17.643594 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.attention_extra_logit : NoneType
I0605 17:37:17.643625 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.attention_mask_summary : False
I0605 17:37:17.643655 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.cast_rotary_position_emb : True
I0605 17:37:17.643687 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.cls : type/praxis.layers.attentions/DotProductAttention
I0605 17:37:17.643719 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combine_qkv : True
I0605 17:37:17.643751 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.643784 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.attention_combine_dims : False
I0605 17:37:17.643817 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.cls : type/praxis.layers.attentions/CombinedQKVProjectionLayer
I0605 17:37:17.643851 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.643883 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.643916 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.dim_per_head : 0
I0605 17:37:17.643948 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.643980 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.644013 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.644045 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.644078 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.644111 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.644143 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.644176 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.644207 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.644239 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.644272 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.644304 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.644336 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.644369 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.644401 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.644433 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.644466 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.explicit_fan_in_fan_out_axes : False
I0605 17:37:17.644498 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.fprop_dtype : NoneType
I0605 17:37:17.644530 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.644562 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.input_dim : 0
I0605 17:37:17.644595 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.mesh_axis_names : NoneType
I0605 17:37:17.644627 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.name : NoneType
I0605 17:37:17.644659 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.num_heads : 0
I0605 17:37:17.644696 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.644729 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.params_init.method : 'xavier'
I0605 17:37:17.644761 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.params_init.scale : 1.000001
I0605 17:37:17.644793 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.644824 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.644855 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.use_bias : True
I0605 17:37:17.644885 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.644916 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.644946 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.644976 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dconv_kernel_size : 3
I0605 17:37:17.645007 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dconv_qkv : False
I0605 17:37:17.645037 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.decode_cache : True
I0605 17:37:17.645068 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dim_per_head : NoneType
I0605 17:37:17.645099 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.645147 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.645180 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.645210 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.645241 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.dropout_at_eval : False
I0605 17:37:17.645272 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.645303 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.645333 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.645364 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.keep_prob : 1.0
I0605 17:37:17.645395 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.645425 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.name : NoneType
I0605 17:37:17.645455 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.noise_shape : NoneType
I0605 17:37:17.645486 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.645516 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.645547 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.645577 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.645608 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.645639 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.645669 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.transpose_qk : False
I0605 17:37:17.645700 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.645731 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.645761 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.fprop_dtype : NoneType
I0605 17:37:17.645791 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.hidden_dim : 0
I0605 17:37:17.645821 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.645852 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.input_dim : 0
I0605 17:37:17.645882 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.internal_enable_per_dim_scale : True
I0605 17:37:17.645913 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.internal_enable_query_scale : True
I0605 17:37:17.645943 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.internal_gshard_gaussian_init : False
I0605 17:37:17.645973 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.mesh_axis_names : NoneType
I0605 17:37:17.646003 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.name : NoneType
I0605 17:37:17.646034 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.ngrammer_tpl : NoneType
I0605 17:37:17.646064 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.num_heads : 1
I0605 17:37:17.646094 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.output_proj_use_nhd_shape : False
I0605 17:37:17.646125 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.646155 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.params_init.method : 'xavier'
I0605 17:37:17.646185 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.params_init.scale : 1.000001
I0605 17:37:17.646216 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.646246 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.attention_combine_dims : False
I0605 17:37:17.646277 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.cls : type/praxis.layers.attentions/AttentionProjection
I0605 17:37:17.646307 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.646338 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.646369 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.dim_per_head : 0
I0605 17:37:17.646399 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.646430 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.646461 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.646491 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.646521 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.646552 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.646582 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.646613 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.646643 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.646673 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.646703 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.646734 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.646764 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.646795 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.646825 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.646855 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.646886 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.explicit_fan_in_fan_out_axes : False
I0605 17:37:17.646916 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.fprop_dtype : NoneType
I0605 17:37:17.646946 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.646977 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.input_dim : 0
I0605 17:37:17.647007 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.is_output_projection : False
I0605 17:37:17.647037 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.mesh_axis_names : NoneType
I0605 17:37:17.647068 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.name : NoneType
I0605 17:37:17.647098 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.num_heads : 0
I0605 17:37:17.647128 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.647159 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.params_init.method : 'xavier'
I0605 17:37:17.647189 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.params_init.scale : 1.000001
I0605 17:37:17.647220 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.647250 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.647281 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.use_bias : True
I0605 17:37:17.647312 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.use_nhd_shape : False
I0605 17:37:17.647342 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.647373 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.647404 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.647434 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.647465 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.647496 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.647526 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.647557 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.647588 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.647619 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.name : NoneType
I0605 17:37:17.647651 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.647684 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.647715 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.647746 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.647777 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.647808 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.pv_einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.647838 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.647868 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.647899 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.647929 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.647960 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.647990 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.648021 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.648051 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.648082 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.name : NoneType
I0605 17:37:17.648113 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.648143 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.648173 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.648204 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.648234 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.648265 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.qk_einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.648295 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.relative_bias_tpl : NoneType
I0605 17:37:17.648326 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.648357 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.cast_as_fprop_dtype : True
I0605 17:37:17.648388 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.cls : type/praxis.layers.embedding_softmax/RotaryPositionalEmbedding
I0605 17:37:17.648419 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.648449 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.648479 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.648510 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.embedding_dims : 0
I0605 17:37:17.648540 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.fprop_dtype : NoneType
I0605 17:37:17.648571 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.648601 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.max_timescale : 10000
I0605 17:37:17.648632 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.mesh_axis_names : NoneType
I0605 17:37:17.648663 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.min_timescale : 1
I0605 17:37:17.648693 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.name : NoneType
I0605 17:37:17.648724 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.648755 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.params_init.method : 'xavier'
I0605 17:37:17.648786 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.params_init.scale : 1.000001
I0605 17:37:17.648816 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.648847 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.648877 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.rotary_position_emb_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.648907 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.scale_logits_by_head_dims : False
I0605 17:37:17.648937 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.648968 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.648998 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.use_bias : False
I0605 17:37:17.649029 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.use_rotary_position_emb : False
I0605 17:37:17.649059 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.weight_split_dims_mapping.dconv : NoneType
I0605 17:37:17.649090 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.weight_split_dims_mapping.proj : NoneType
I0605 17:37:17.649127 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.649159 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_atten_tpl.zero_fully_masked : False
I0605 17:37:17.649190 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_split_dims_mapping.ffn0 : NoneType
I0605 17:37:17.649220 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_split_dims_mapping.ffn1 : NoneType
I0605 17:37:17.649250 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.649281 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.649312 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.approximate : True
I0605 17:37:17.649342 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.cls : type/praxis.layers.activations/GELU
I0605 17:37:17.649373 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.649403 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.649433 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.649464 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.fprop_dtype : NoneType
I0605 17:37:17.649494 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.649525 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.mesh_axis_names : NoneType
I0605 17:37:17.649555 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.name : NoneType
I0605 17:37:17.649586 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.649616 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.params_init.method : 'xavier'
I0605 17:37:17.649647 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.params_init.scale : 1.000001
I0605 17:37:17.649677 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.649708 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.649739 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.activation_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.649769 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.add_skip_connection : True
I0605 17:37:17.649799 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.apply_padding_first : False
I0605 17:37:17.649830 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.cls : type/praxis.layers.transformers/TransformerFeedForward
I0605 17:37:17.649861 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.649891 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.649922 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.649952 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.649982 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.650013 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.cls : type/praxis.layers.activations/ReLU
I0605 17:37:17.650044 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.650074 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.650104 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.650135 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.fprop_dtype : NoneType
I0605 17:37:17.650165 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.650195 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.mesh_axis_names : NoneType
I0605 17:37:17.650226 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.name : NoneType
I0605 17:37:17.650256 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.650288 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.params_init.method : 'xavier'
I0605 17:37:17.650318 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.params_init.scale : 1.000001
I0605 17:37:17.650348 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.650379 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.650409 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.activation_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.650440 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.bias_init : 0.0
I0605 17:37:17.650470 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.cls : type/praxis.layers.linears/FeedForward
I0605 17:37:17.650501 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.650531 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.650561 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.650592 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.fprop_dtype : NoneType
I0605 17:37:17.650622 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.has_bias : True
I0605 17:37:17.650655 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.650688 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.input_dims : 0
I0605 17:37:17.650718 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.650749 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.cls : type/praxis.layers.linears/Linear
I0605 17:37:17.650779 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.650809 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.650840 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.650871 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.650901 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.cls : type/praxis.layers.base_ops/Einsum
I0605 17:37:17.650932 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.650963 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.650994 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.651025 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.fprop_dtype : NoneType
I0605 17:37:17.651055 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.651086 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.mesh_axis_names : NoneType
I0605 17:37:17.651116 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.name : NoneType
I0605 17:37:17.651147 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.651178 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.params_init.method : 'xavier'
I0605 17:37:17.651209 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.params_init.scale : 1.000001
I0605 17:37:17.651240 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.651270 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.651301 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.einsum_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.651332 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.fprop_dtype : NoneType
I0605 17:37:17.651362 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.651392 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.input_dims : 0
I0605 17:37:17.651423 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.mesh_axis_names : NoneType
I0605 17:37:17.651454 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.name : NoneType
I0605 17:37:17.651484 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.output_dims : 0
I0605 17:37:17.651514 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.651545 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.params_init.method : 'xavier'
I0605 17:37:17.651576 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.params_init.scale : 1.000001
I0605 17:37:17.651606 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.651637 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.651667 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.weight_init : NoneType
I0605 17:37:17.651697 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.651728 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.mesh_axis_names : NoneType
I0605 17:37:17.651758 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.name : NoneType
I0605 17:37:17.651789 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.output_dims : 0
I0605 17:37:17.651820 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.651850 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.params_init.method : 'xavier'
I0605 17:37:17.651881 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.params_init.scale : 1.000001
I0605 17:37:17.651912 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.651942 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.651973 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.weight_init : NoneType
I0605 17:37:17.652003 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.652033 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.fprop_dtype : NoneType
I0605 17:37:17.652064 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.has_bias : True
I0605 17:37:17.652095 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.hidden_dims : 0
I0605 17:37:17.652125 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.652156 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.input_dims : 0
I0605 17:37:17.652187 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.internal_gshard_variance_scaling_fan_in_init : False
I0605 17:37:17.652217 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.652247 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.cls : type/praxis.layers.normalizations/LayerNorm
I0605 17:37:17.652277 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.652308 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.652339 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.dim : 0
I0605 17:37:17.652369 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.652400 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.epsilon : 1e-06
I0605 17:37:17.652431 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.fprop_dtype : NoneType
I0605 17:37:17.652461 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.652492 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.mesh_axis_names : NoneType
I0605 17:37:17.652522 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.name : NoneType
I0605 17:37:17.652552 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.652583 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.params_init.method : 'xavier'
I0605 17:37:17.652613 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.params_init.scale : 1.000001
I0605 17:37:17.652643 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.reductions_in_fp32 : False
I0605 17:37:17.652673 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.652704 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.652735 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.use_bias : True
I0605 17:37:17.652765 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.use_scale : True
I0605 17:37:17.652795 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.ln_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.652826 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.mesh_axis_names : NoneType
I0605 17:37:17.652857 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.name : NoneType
I0605 17:37:17.652887 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.norm_policy : 'pre'
I0605 17:37:17.652917 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.output_dims : 0
I0605 17:37:17.652948 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.652978 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.params_init.method : 'xavier'
I0605 17:37:17.653009 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.params_init.scale : 1.000001
I0605 17:37:17.653039 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_prob : 0.0
I0605 17:37:17.653069 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.653105 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.653138 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.653170 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.653200 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.dropout_at_eval : False
I0605 17:37:17.653231 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.653262 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.653292 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.653323 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.keep_prob : 1.0
I0605 17:37:17.653354 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.653384 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.name : NoneType
I0605 17:37:17.653415 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.noise_shape : NoneType
I0605 17:37:17.653446 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.653477 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.653507 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.653537 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.653568 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.653599 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.653630 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.transpose_qk : False
I0605 17:37:17.653660 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.relu_dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.653691 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_prob : 0.0
I0605 17:37:17.653725 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.activation_split_dims_mapping.out : NoneType
I0605 17:37:17.653755 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.cls : type/praxis.layers.stochastics/Dropout
I0605 17:37:17.653786 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.653817 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.653847 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.dropout_at_eval : False
I0605 17:37:17.653878 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.653909 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.fprop_dtype : NoneType
I0605 17:37:17.653939 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.653970 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.keep_prob : 1.0
I0605 17:37:17.654000 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.mesh_axis_names : NoneType
I0605 17:37:17.654031 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.name : NoneType
I0605 17:37:17.654061 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.noise_shape : NoneType
I0605 17:37:17.654092 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.noise_shape_broadcast_dims : NoneType
I0605 17:37:17.654122 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.654153 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.params_init.method : 'xavier'
I0605 17:37:17.654186 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.params_init.scale : 1.000001
I0605 17:37:17.654217 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.654248 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.654278 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.transpose_qk : False
I0605 17:37:17.654308 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_dropout_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.654339 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_droppath_prob : 0.0
I0605 17:37:17.654369 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.residual_weight : 1.0
I0605 17:37:17.654399 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.654429 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.654459 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.use_gated_activation : False
I0605 17:37:17.654489 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.weight_split_dims_mapping.ffn0 : NoneType
I0605 17:37:17.654520 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.weight_split_dims_mapping.ffn1 : NoneType
I0605 17:37:17.654550 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.tr_fflayer_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.654581 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.use_cross_attention : False
I0605 17:37:17.654611 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.transformer_layer_params_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.654642 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.unadjusted_expert_capacity_factor : 2.0
I0605 17:37:17.654674 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.use_cross_attention : False
I0605 17:37:17.654704 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.block.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.654735 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.checkpoint_policy : 'save_nothing'
I0605 17:37:17.654766 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.cls : type/praxis.layers.transformers/StackedTransformerRepeated
I0605 17:37:17.654796 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.contiguous_submeshes : NoneType
I0605 17:37:17.654826 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.dcn_mesh_shape : NoneType
I0605 17:37:17.654856 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.dtype : type/jax.numpy/float32
I0605 17:37:17.654886 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.fprop_dtype : NoneType
I0605 17:37:17.654917 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.ici_mesh_shape : NoneType
I0605 17:37:17.654947 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.mesh_axis_names : NoneType
I0605 17:37:17.654977 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.name : NoneType
I0605 17:37:17.655007 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.nd_prefix_shape : NoneType
I0605 17:37:17.655038 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.params_init.cls : type/praxis.base_layer/WeightInit
I0605 17:37:17.655068 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.params_init.method : 'xavier'
I0605 17:37:17.655098 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.params_init.scale : 1.000001
I0605 17:37:17.655129 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.repeat_layer_name : 'repeat'
I0605 17:37:17.655158 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.repeat_optimizer_dims_mapping : NoneType
I0605 17:37:17.655188 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.shared_weight_layer_id : NoneType
I0605 17:37:17.655219 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.skip_lp_regularization : NoneType
I0605 17:37:17.655249 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.sublayer_name : 'sub'
I0605 17:37:17.655279 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.unroll_in_decode : True
I0605 17:37:17.655310 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.weight_split_dims_mapping.block : NoneType
I0605 17:37:17.655340 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.655370 139756486557120 train.py:171]   model.lm_tpl.stacked_transformer_tpl.x_times : 24
I0605 17:37:17.655400 139756486557120 train.py:171]   model.lm_tpl.vocab_size : 51200
I0605 17:37:17.655430 139756486557120 train.py:171]   model.lm_tpl.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.655461 139756486557120 train.py:171]   model.mesh_axis_names : NoneType
I0605 17:37:17.655491 139756486557120 train.py:171]   model.model_type : 'causal'
I0605 17:37:17.655521 139756486557120 train.py:171]   model.name : 'xformer_lm'
I0605 17:37:17.655550 139756486557120 train.py:171]   model.params_init.method : 'gaussian'
I0605 17:37:17.655580 139756486557120 train.py:171]   model.params_init.scale : 0.023
I0605 17:37:17.655610 139756486557120 train.py:171]   model.report_strict_acc : False
I0605 17:37:17.655640 139756486557120 train.py:171]   model.return_predictions : False
I0605 17:37:17.655670 139756486557120 train.py:171]   model.shared_weight_layer_id : NoneType
I0605 17:37:17.655700 139756486557120 train.py:171]   model.skip_lp_regularization : NoneType
I0605 17:37:17.655730 139756486557120 train.py:171]   model.weight_split_dims_mapping.wt : NoneType
I0605 17:37:17.655761 139756486557120 train.py:171]   name : 'xformer_task'
I0605 17:37:17.655791 139756486557120 train.py:171]   summary_verbosity : 3
I0605 17:37:17.655821 139756486557120 train.py:171]   train.always_use_train_for_model_init : True
I0605 17:37:17.655851 139756486557120 train.py:171]   train.apply_mutable_list : ['aux_loss', 'summaries', 'non_trainable', 'batch_stats', 'params_axes']
I0605 17:37:17.655882 139756486557120 train.py:171]   train.async_summary_writing : True
I0605 17:37:17.655912 139756486557120 train.py:171]   train.cls : type/paxml.tasks_lib/SingleTask.Train
I0605 17:37:17.655942 139756486557120 train.py:171]   train.decode_interval_steps : NoneType
I0605 17:37:17.655972 139756486557120 train.py:171]   train.decode_start_after_n_steps : 0
I0605 17:37:17.656003 139756486557120 train.py:171]   train.decode_use_ema_states : False
I0605 17:37:17.656033 139756486557120 train.py:171]   train.device_sync_interval_steps : NoneType
I0605 17:37:17.656064 139756486557120 train.py:171]   train.enable_input_checkpointing : False
I0605 17:37:17.656093 139756486557120 train.py:171]   train.enforce_input_specs : False
I0605 17:37:17.656123 139756486557120 train.py:171]   train.eval_interval_steps : 100
I0605 17:37:17.656153 139756486557120 train.py:171]   train.eval_skip_train : False
I0605 17:37:17.656183 139756486557120 train.py:171]   train.eval_use_ema_states : False
I0605 17:37:17.656213 139756486557120 train.py:171]   train.external_checkpoint_handler : NoneType
I0605 17:37:17.656244 139756486557120 train.py:171]   train.external_checkpoint_path : NoneType
I0605 17:37:17.656274 139756486557120 train.py:171]   train.inputs_split_mapping : NoneType
I0605 17:37:17.656304 139756486557120 train.py:171]   train.learner.check_valid_step : True
I0605 17:37:17.656334 139756486557120 train.py:171]   train.learner.cls : type/paxml.learners/Learner
I0605 17:37:17.656365 139756486557120 train.py:171]   train.learner.enable_skip_step_on_gradient_anomalies : True
I0605 17:37:17.656395 139756486557120 train.py:171]   train.learner.force_repeat_prefix_structure : False
I0605 17:37:17.656425 139756486557120 train.py:171]   train.learner.grad_norm_individual_vars : False
I0605 17:37:17.656456 139756486557120 train.py:171]   train.learner.grad_norm_summary : True
I0605 17:37:17.656486 139756486557120 train.py:171]   train.learner.keep_optimizer_state_for_excluded_vars : False
I0605 17:37:17.656517 139756486557120 train.py:171]   train.learner.loss_name : 'total_loss'
I0605 17:37:17.656547 139756486557120 train.py:171]   train.learner.name : ''
I0605 17:37:17.656577 139756486557120 train.py:171]   train.learner.optimizer.beta1 : 0.9
I0605 17:37:17.656607 139756486557120 train.py:171]   train.learner.optimizer.beta2 : 0.95
I0605 17:37:17.656638 139756486557120 train.py:171]   train.learner.optimizer.clip_gradient_norm_to_value : 1.0
I0605 17:37:17.656668 139756486557120 train.py:171]   train.learner.optimizer.clip_gradient_single_norm_to_value : 0.0
I0605 17:37:17.656698 139756486557120 train.py:171]   train.learner.optimizer.clip_threshold : 1.0
I0605 17:37:17.656729 139756486557120 train.py:171]   train.learner.optimizer.cls : type/praxis.optimizers/Adam
I0605 17:37:17.656759 139756486557120 train.py:171]   train.learner.optimizer.decoupled_weight_decay : NoneType
I0605 17:37:17.656790 139756486557120 train.py:171]   train.learner.optimizer.ema_decay : 0.0
I0605 17:37:17.656820 139756486557120 train.py:171]   train.learner.optimizer.epsilon : 1e-08
I0605 17:37:17.656850 139756486557120 train.py:171]   train.learner.optimizer.epsilon_root : 0.0
I0605 17:37:17.656880 139756486557120 train.py:171]   train.learner.optimizer.ewc_regularizer_weight : 0.0
I0605 17:37:17.656910 139756486557120 train.py:171]   train.learner.optimizer.ewc_weight_per_var : NoneType
I0605 17:37:17.656940 139756486557120 train.py:171]   train.learner.optimizer.l1_regularizer_weight : NoneType
I0605 17:37:17.656970 139756486557120 train.py:171]   train.learner.optimizer.l2_regularizer_weight : NoneType
I0605 17:37:17.657001 139756486557120 train.py:171]   train.learner.optimizer.learning_rate : 0.0006
I0605 17:37:17.657031 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.cls : type/praxis.schedules/LinearRampupCosineDecay
I0605 17:37:17.657062 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.decay_end : 500000
I0605 17:37:17.657092 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.decay_start : 1
I0605 17:37:17.657137 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.max : 1.0
I0605 17:37:17.657169 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.min_ratio : 0.1
I0605 17:37:17.657199 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.name : ''
I0605 17:37:17.657230 139756486557120 train.py:171]   train.learner.optimizer.lr_schedule.warmup_steps : 0
I0605 17:37:17.657260 139756486557120 train.py:171]   train.learner.optimizer.maybe_inf_to_nan : True
I0605 17:37:17.657290 139756486557120 train.py:171]   train.learner.optimizer.name : ''
I0605 17:37:17.657320 139756486557120 train.py:171]   train.learner.optimizer.sharded_adam : True
I0605 17:37:17.657351 139756486557120 train.py:171]   train.learner.optimizer.skip_lp_1d_vectors : False
I0605 17:37:17.657380 139756486557120 train.py:171]   train.learner.optimizer.weight_decay : 0.001
I0605 17:37:17.657411 139756486557120 train.py:171]   train.learner.repeat_prefix_sep : '#'
I0605 17:37:17.657441 139756486557120 train.py:171]   train.learner.skip_step_gradient_norm_value : 0.0
I0605 17:37:17.657471 139756486557120 train.py:171]   train.learner.skip_zero_gradients : NoneType
I0605 17:37:17.657502 139756486557120 train.py:171]   train.learner.stochastic_gradient : NoneType
I0605 17:37:17.657532 139756486557120 train.py:171]   train.learner.var_norm_summary : True
I0605 17:37:17.657562 139756486557120 train.py:171]   train.learner.vectorize_on_repeat_prefix : True
I0605 17:37:17.657592 139756486557120 train.py:171]   train.log_train_output_interval_steps : NoneType
I0605 17:37:17.657623 139756486557120 train.py:171]   train.max_inflight_steps : 2
I0605 17:37:17.657653 139756486557120 train.py:171]   train.num_train_steps : 10000000.0
I0605 17:37:17.657683 139756486557120 train.py:171]   train.profiler_capture_step : NoneType
I0605 17:37:17.657713 139756486557120 train.py:171]   train.profiler_max_num_hosts : NoneType
I0605 17:37:17.657743 139756486557120 train.py:171]   train.profiler_min_duration_sec : 1
I0605 17:37:17.657773 139756486557120 train.py:171]   train.profiler_num_steps : 2
I0605 17:37:17.657804 139756486557120 train.py:171]   train.random_seed : 1234
I0605 17:37:17.657834 139756486557120 train.py:171]   train.restore_transformations : NoneType
I0605 17:37:17.657865 139756486557120 train.py:171]   train.save_interval_steps : 100000
I0605 17:37:17.657895 139756486557120 train.py:171]   train.save_keep_interval_duration : '12h'
I0605 17:37:17.657925 139756486557120 train.py:171]   train.save_max_to_keep : 10
I0605 17:37:17.657955 139756486557120 train.py:171]   train.summary_accumulate_interval_steps : NoneType
I0605 17:37:17.657985 139756486557120 train.py:171]   train.summary_interval_steps : 100
I0605 17:37:17.658015 139756486557120 train.py:171]   train.tensorstore_metadata_key : NoneType
I0605 17:37:17.658048 139756486557120 train.py:171]   train.variable_norm_summary : True
I0605 17:37:17.658079 139756486557120 train.py:171]   vn.cls : type/paxml.tasks_lib/SingleTask.VariationalNoise
I0605 17:37:17.658113 139756486557120 train.py:171]   vn.vn_regex : ''
I0605 17:37:17.658157 139756486557120 train.py:171]   vn.vn_scale : 0.0
I0605 17:37:17.658191 139756486557120 train.py:171]   vn.vn_start_step : 0
I0605 17:37:17.658233 139756486557120 train.py:173] [PAX STATUS]: Initializing decoder
I0605 17:37:17.658369 139756486557120 checkpoint_creators.py:564] [PAX STATUS]: Creating checkpointer.
I0605 17:37:17.658570 139756486557120 py_utils.py:338] Starting sync_global_devices checkpointer:makedirs:log_NVIDIA1_3BPmap/checkpoints across 1 devices globally
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0000d1b00 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:0}, signal={0x6060001682c0:1} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:1}, signal={0x6060001682c0:2} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000167a80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000167960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000167a80, semaphore=0x6060001682c0, value=2 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000167960, semaphore=0x6060001682c0, value=2 (OK)
W0605 17:37:17.731722 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0006318092346191406 sec
W0605 17:37:17.732191 139756486557120 dispatch.py:272] Finished tracing + transforming _psum for pjit in 0.0015490055084228516 sec
W0605 17:37:17.733086 139756486557120 pxla.py:1882] Compiling _psum for with global shapes and types [ShapedArray(uint32[1])]. Argument mapping: (GSPMDSharding({maximal device=0}),).
W0605 17:37:17.738058 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_psum) in 0.004811763763427734 sec
W0605 17:37:17.966061 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_psum) in 0.22601723670959473 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000203720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000203720, semaphore=0x6060001683e0, value=0 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=1, fence=0x60400081fad0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000203720, from_fence=0x606000167a80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000167960, semaphore=0x6060001683e0, value=1 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x606000203720 {0x6060001683e0:0, 0x6060001682c0:2}, signal_fence=0x60400081fad0 {0x6060001683e0:1} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000220ee0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000221300 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000220ee0, semaphore=0x6060001683e0, value=1 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000221300, semaphore=0x6060001683e0, value=1 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc14a40 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:1}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0000d1c80, wait={0x6060001682c0:2, 0x6060001683e0:1}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0000d1c80, wait={0x6060001683e0:1}, signal={} (OK)
I0605 17:37:17.968034 139756486557120 py_utils.py:341] Finished sync_global_devices checkpointer:makedirs:log_NVIDIA1_3BPmap/checkpoints across 1 devices globally
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0000a5a00 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:2}, signal={0x6060001682c0:3} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:3}, signal={0x6060001682c0:4} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cfe40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cfcc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cfe40, semaphore=0x6060001682c0, value=4 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cfcc0, semaphore=0x6060001682c0, value=4 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cec40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cec40, semaphore=0x6060001683e0, value=1 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=2, fence=0x60400081f1d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060002cec40, from_fence=0x6060002cfe40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cfcc0, semaphore=0x6060001683e0, value=2 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060002cec40 {0x6060001683e0:1, 0x6060001682c0:4}, signal_fence=0x60400081f1d0 {0x6060001683e0:2} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002ceb80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002ceac0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002ceb80, semaphore=0x6060001683e0, value=2 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002ceac0, semaphore=0x6060001683e0, value=2 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc14300 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:2}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0000a5ac0, wait={0x6060001682c0:4, 0x6060001683e0:2}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0000a5ac0, wait={0x6060001683e0:2}, signal={} (OK)
I0605 17:37:17.971108 139756486557120 utils.py:366] Cleaning up existing temporary directories at log_NVIDIA1_3BPmap/checkpoints.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00010e280 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:4}, signal={0x6060001682c0:5} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:5}, signal={0x6060001682c0:6} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cd9e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cdb00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cd9e0, semaphore=0x6060001682c0, value=6 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cdb00, semaphore=0x6060001682c0, value=6 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002d3d40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d3d40, semaphore=0x6060001683e0, value=2 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=3, fence=0x60400081e990 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060002d3d40, from_fence=0x6060002cd9e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cdb00, semaphore=0x6060001683e0, value=3 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060002d3d40 {0x6060001683e0:2, 0x6060001682c0:6}, signal_fence=0x60400081e990 {0x6060001683e0:3} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002d3ec0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002d3f20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d3ec0, semaphore=0x6060001683e0, value=3 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d3f20, semaphore=0x6060001683e0, value=3 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13f60 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:3}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00010df80, wait={0x6060001682c0:6, 0x6060001683e0:3}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00010df80, wait={0x6060001683e0:3}, signal={} (OK)
I0605 17:37:17.973646 139756486557120 train.py:206] [PAX STATUS]: Creating task
I0605 17:37:18.160361 139756486557120 train.py:217] [PAX STATUS]: Initializing partitioner
I0605 17:37:18.160518 139756486557120 partitioning.py:576] Using pmap for data parallelism.
I0605 17:37:18.160575 139756486557120 train.py:245] [PAX STATUS]: Creating executor.
I0605 17:37:18.160630 139756486557120 train.py:249] [PAX STATUS]: Setting up executor.
W0605 17:37:18.164295 139756486557120 dispatch.py:272] Finished tracing + transforming jit(convert_element_type) in 0.0002434253692626953 sec
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00010d740 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:6}, signal={0x6060001682c0:7} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:7}, signal={0x6060001682c0:8} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002d30e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002d3080 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d30e0, semaphore=0x6060001682c0, value=8 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d3080, semaphore=0x6060001682c0, value=8 (OK)
W0605 17:37:18.166568 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00041937828063964844 sec
W0605 17:37:18.167441 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_seed for pjit in 0.0018720626831054688 sec
W0605 17:37:18.168067 139756486557120 pxla.py:1882] Compiling _threefry_seed for with global shapes and types [ShapedArray(int32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:18.172546 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_seed) in 0.004347562789916992 sec
W0605 17:37:18.443973 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_seed) in 0.2711031436920166 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000a70c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a70c0, semaphore=0x6060001683e0, value=3 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=4, fence=0x60400055c890 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060000a70c0, from_fence=0x6060002d30e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002d3080, semaphore=0x6060001683e0, value=4 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00053f590, f=0, wait_fence=0x6060000a70c0 {0x6060001683e0:3, 0x6060001682c0:8}, signal_fence=0x60400055c890 {0x6060001683e0:4} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000a6f40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000a6ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a6f40, semaphore=0x6060001683e0, value=4 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a6ee0, semaphore=0x6060001683e0, value=4 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00010d500, wait={0x6060001682c0:8, 0x6060001683e0:4}, signal={} (OK)
I0605 17:37:18.445732 139756486557120 partitioning.py:420] input_p.tf_data_service_address: None
I0605 17:37:18.445963 139756486557120 executors.py:163] [PAX STATUS]: Instantiating train input pipeline.
I0605 17:37:18.449376 139756486557120 executors.py:222] [PAX STATUS]: Setting up partitioner
I0605 17:37:18.449437 139756486557120 partitioning.py:353] [PAX STATUS]: Getting input shapes from first batch.
I0605 17:37:19.157606 139756486557120 local.py:50] Created artifact Input specs of type ArtifactType.FILE and value log_NVIDIA1_3BPmap/input_specs.json.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000164ec0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:8}, signal={0x6060001682c0:9} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:9}, signal={0x6060001682c0:10} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600040c5e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600040c640 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c5e0, semaphore=0x6060001682c0, value=10 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c640, semaphore=0x6060001682c0, value=10 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600040c700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c700, semaphore=0x6060001683e0, value=4 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=5, fence=0x6040002c8ad0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600040c700, from_fence=0x60600040c5e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c640, semaphore=0x6060001683e0, value=5 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00053f590, f=0, wait_fence=0x60600040c700 {0x6060001683e0:4, 0x6060001682c0:10}, signal_fence=0x6040002c8ad0 {0x6060001683e0:5} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600040c820 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600040c880 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c820, semaphore=0x6060001683e0, value=5 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600040c880, semaphore=0x6060001683e0, value=5 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000164f80, wait={0x6060001682c0:10, 0x6060001683e0:5}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000165280, wait={0x6060001683e0:5}, signal={} (OK)
W0605 17:37:19.679517 139756486557120 optimizers.py:1170] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
I0605 17:37:19.679631 139756486557120 optimizers.py:1173] Using sharded_adam.
W0605 17:37:19.679672 139756486557120 optimizers.py:580] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
W0605 17:37:19.693978 139756486557120 optimizers.py:1170] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
I0605 17:37:19.694039 139756486557120 optimizers.py:1173] Using sharded_adam.
W0605 17:37:19.694076 139756486557120 optimizers.py:580] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
I0605 17:37:19.708143 139756486557120 trainer_lib.py:197] post_init_model_params: log_NVIDIA1_3BPmap/post_init_model_params.txt
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000177c40 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:10}, signal={0x6060001682c0:11} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:11}, signal={0x6060001682c0:12} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600046c340 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600046c3a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c340, semaphore=0x6060001682c0, value=12 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c3a0, semaphore=0x6060001682c0, value=12 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600046c460 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c460, semaphore=0x6060001683e0, value=5 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=6, fence=0x6040002b6590 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600046c460, from_fence=0x60600046c340 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c3a0, semaphore=0x6060001683e0, value=6 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00053f590, f=0, wait_fence=0x60600046c460 {0x6060001683e0:5, 0x6060001682c0:12}, signal_fence=0x6040002b6590 {0x6060001683e0:6} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600046c580 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600046c5e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c580, semaphore=0x6060001683e0, value=6 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600046c5e0, semaphore=0x6060001683e0, value=6 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000177d00, wait={0x6060001682c0:12, 0x6060001683e0:6}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000178000, wait={0x6060001683e0:6}, signal={} (OK)
W0605 17:37:19.902000 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00021767616271972656 sec
W0605 17:37:19.902926 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0017635822296142578 sec
W0605 17:37:19.904009 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.0032732486724853516 sec
W0605 17:37:19.904708 139756486557120 pxla.py:1882] Compiling _threefry_split_original for with global shapes and types [ShapedArray(uint32[2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:19.909348 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0005061626434326172 sec
W0605 17:37:19.910399 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003032684326171875 sec
W0605 17:37:19.911373 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00029921531677246094 sec
W0605 17:37:19.912310 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003139972686767578 sec
W0605 17:37:19.912992 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002911090850830078 sec
W0605 17:37:19.959887 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_split_original) in 0.05504131317138672 sec
W0605 17:37:21.409307 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_split_original) in 1.4490509033203125 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000258fe0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000258fe0, semaphore=0x6060001683e0, value=6 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=7, fence=0x604000e5abd0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000258fe0, from_fence=0x6060000a6f40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a6ee0, semaphore=0x6060001683e0, value=7 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00047e130, f=0, wait_fence=0x606000258fe0 {0x6060001683e0:6}, signal_fence=0x604000e5abd0 {0x6060001683e0:7} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000258ec0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002590a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000258ec0, semaphore=0x6060001683e0, value=7 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002590a0, semaphore=0x6060001683e0, value=7 (OK)
W0605 17:37:21.413856 139756486557120 dispatch.py:272] Finished tracing + transforming _unstack for pjit in 0.0012726783752441406 sec
W0605 17:37:21.415073 139756486557120 pxla.py:1882] Compiling _unstack for with global shapes and types [ShapedArray(uint32[2,2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:21.421748 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_unstack) in 0.0064165592193603516 sec
W0605 17:37:21.630789 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_unstack) in 0.20858287811279297 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000328d00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000328d00, semaphore=0x6060001683e0, value=7 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=8, fence=0x604000cef690 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000328d00, from_fence=0x606000258ec0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002590a0, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000314e10, f=0, wait_fence=0x606000328d00 {0x6060001683e0:7}, signal_fence=0x604000cef690 {0x6060001683e0:8} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000328ee0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003292a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000328ee0, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003292a0, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000328e20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000329300 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000328e20, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000329300, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00035a100, wait={0x6060001683e0:8}, signal={} (OK)
W0605 17:37:21.633492 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00022339820861816406 sec
W0605 17:37:21.634345 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.001689910888671875 sec
W0605 17:37:21.635529 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.00327301025390625 sec
W0605 17:37:21.636210 139756486557120 pxla.py:1882] Compiling _threefry_split_original for with global shapes and types [ShapedArray(uint32[2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:21.641349 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.000331878662109375 sec
W0605 17:37:21.642361 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003070831298828125 sec
W0605 17:37:21.643281 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003108978271484375 sec
W0605 17:37:21.644005 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00032782554626464844 sec
W0605 17:37:21.691162 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_split_original) in 0.054818153381347656 sec
W0605 17:37:23.174265 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_split_original) in 1.4827439785003662 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a72a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a72a0, semaphore=0x6060001683e0, value=8 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=9, fence=0x6040008b7650 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003a72a0, from_fence=0x606000328e20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000329300, semaphore=0x6060001683e0, value=9 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0000ec8f0, f=0, wait_fence=0x6060003a72a0 {0x6060001683e0:8}, signal_fence=0x6040008b7650 {0x6060001683e0:9} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a6940 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a7900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a6940, semaphore=0x6060001683e0, value=9 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a7900, semaphore=0x6060001683e0, value=9 (OK)
W0605 17:37:23.179615 139756486557120 dispatch.py:272] Finished tracing + transforming _unstack for pjit in 0.0019333362579345703 sec
W0605 17:37:23.181048 139756486557120 pxla.py:1882] Compiling _unstack for with global shapes and types [ShapedArray(uint32[4,2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:23.186113 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_unstack) in 0.004763603210449219 sec
W0605 17:37:23.414857 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_unstack) in 0.22841119766235352 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005df200 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005df200, semaphore=0x6060001683e0, value=9 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=10, fence=0x60400025bf90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005df200, from_fence=0x6060003a6940 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a7900, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0002d79b0, f=0, wait_fence=0x6060005df200 {0x6060001683e0:9}, signal_fence=0x60400025bf90 {0x6060001683e0:10} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000193d60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001943c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193d60, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001943c0, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000195620 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000193460 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000195620, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193460, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001933a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000193340 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001933a0, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193340, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000193b20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000193c40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193b20, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193c40, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0002e7bc0, wait={0x6060001683e0:10}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13ee0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:10}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13ee0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:10}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13ee0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:10}, signal={} (OK)
I0605 17:37:23.416016 139756486557120 trainer_lib.py:378] init_var prng_seed: {'params': Array([1477712937, 1244108694], dtype=uint32), 'random': Array([713085529, 937672790], dtype=uint32), 'dropout': Array([3893856254, 2733895282], dtype=uint32)}
I0605 17:37:23.417782 139756486557120 trainer_lib.py:379] var_weight_hparams: {'params': {'lm': {'final_ln': {'bias': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=None, repeat_prefix_split_dims_mapping=None, repeat_optimizer_dims_mapping=None, fan_in_axes=None, fan_out_axes=None), 'scale': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=None, repeat_prefix_split_dims_mapping=None, repeat_optimizer_dims_mapping=None, fan_in_axes=None, fan_out_axes=None)}, 'position_emb': {'emb_var': WeightHParams(shape=[2048, 2048], init=WeightInit(method='gaussian', scale=0.023), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=None, repeat_prefix_split_dims_mapping=None, repeat_optimizer_dims_mapping=None, fan_in_axes=None, fan_out_axes=None)}, 'softmax': {'logits_ffn': {'bias': {'b': WeightHParams(shape=[51200], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=None, repeat_prefix_split_dims_mapping=None, repeat_optimizer_dims_mapping=None, fan_in_axes=None, fan_out_axes=None)}, 'linear': {'w': WeightHParams(shape=[2048, 51200], init=WeightInit(method='gaussian', scale=0.022097086912079608), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=None, repeat_prefix_split_dims_mapping=None, repeat_optimizer_dims_mapping=None, fan_in_axes=None, fan_out_axes=None)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': WeightHParams(shape=[8192], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}, 'linear': {'w': WeightHParams(shape=[2048, 8192], init=WeightInit(method='gaussian', scale=0.023), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}}, 'ffn_layer2': {'bias': {'b': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}, 'linear': {'w': WeightHParams(shape=[8192, 2048], init=WeightInit(method='gaussian', scale=0.023), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}}, 'layer_norm': {'bias': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None), 'scale': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}}, 'layer_norm': {'bias': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None), 'scale': WeightHParams(shape=[2048], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=['__lingvo_jax_skip_regularization'], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}, 'self_attention': {'combined_qkv': {'w': WeightHParams(shape=[3, 2048, 32, 64], init=WeightInit(method='gaussian', scale=0.023), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}, 'per_dim_scale': {'per_dim_scale': WeightHParams(shape=[64], init=WeightInit(method='constant', scale=0.0), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}, 'post': {'w': WeightHParams(shape=[2048, 32, 64], init=WeightInit(method='gaussian', scale=0.023), dtype=<class 'jax.numpy.float32'>, collections=[], mesh_shape=None, tensor_split_dims_mapping=None, repeat_prefix=[24], repeat_prefix_split_dims_mapping=(-1,), repeat_optimizer_dims_mapping=(-1,), fan_in_axes=None, fan_out_axes=None)}}}}}}}}}
I0605 17:37:23.466831 139756486557120 base_layer.py:632] Creating var /lm/softmax/logits_ffn/linear/w with shape=[2048, 51200], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.022097086912079608
I0605 17:37:23.472266 139756486557120 base_layer.py:632] Creating var /lm/position_emb/emb_var with shape=[2048, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.561209 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.562514 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.578722 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/combined_qkv/w with shape=[3, 2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.583390 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/per_dim_scale/per_dim_scale with shape=[64], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.593972 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/post/w with shape=[2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.618214 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.619433 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.631588 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/linear/w with shape=[2048, 8192], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.635301 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/bias/b with shape=[8192], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.647808 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/linear/w with shape=[8192, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.651491 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/bias/b with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.706857 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.708141 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.731529 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/combined_qkv/w with shape=[3, 2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.736371 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/per_dim_scale/per_dim_scale with shape=[64], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.746949 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/post/w with shape=[2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.771519 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.772733 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.785490 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/linear/w with shape=[2048, 8192], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.789337 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/bias/b with shape=[8192], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.803040 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/linear/w with shape=[8192, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:37:23.806895 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/bias/b with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.824109 139756486557120 base_layer.py:632] Creating var /lm/final_ln/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.825375 139756486557120 base_layer.py:632] Creating var /lm/final_ln/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:37:23.831791 139756486557120 base_layer.py:632] Creating var /lm/softmax/logits_ffn/bias/b with shape=[51200], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:37:23.844305 139756486557120 dispatch.py:272] Finished tracing + transforming init_fn for pjit in 0.42580294609069824 sec
W0605 17:37:23.849604 139756486557120 pxla.py:1882] Compiling init_fn for with global shapes and types [ShapedArray(uint32[2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:23.854750 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003132820129394531 sec
W0605 17:37:23.855348 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_seed for pjit in 0.0013680458068847656 sec
W0605 17:37:23.856750 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.0001819133758544922 sec
W0605 17:37:23.857589 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0015332698822021484 sec
W0605 17:37:23.858543 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_fold_in for pjit in 0.004767656326293945 sec
W0605 17:37:23.862845 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00031375885009765625 sec
W0605 17:37:23.863823 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003020763397216797 sec
W0605 17:37:23.864712 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.000301361083984375 sec
W0605 17:37:23.865507 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004038810729980469 sec
W0605 17:37:23.915956 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00019168853759765625 sec
W0605 17:37:23.916870 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0016338825225830078 sec
W0605 17:37:23.917910 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.0029790401458740234 sec
W0605 17:37:23.921130 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00031447410583496094 sec
W0605 17:37:23.922111 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003139972686767578 sec
W0605 17:37:23.923000 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002968311309814453 sec
W0605 17:37:23.923702 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003113746643066406 sec
W0605 17:37:24.028707 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00019669532775878906 sec
W0605 17:37:24.029680 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.001708984375 sec
W0605 17:37:24.030701 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.0030384063720703125 sec
W0605 17:37:24.033909 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003120899200439453 sec
W0605 17:37:24.034902 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00031876564025878906 sec
W0605 17:37:24.035800 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00029659271240234375 sec
W0605 17:37:24.036511 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003082752227783203 sec
W0605 17:37:24.088186 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00018739700317382812 sec
W0605 17:37:24.088986 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.001508474349975586 sec
W0605 17:37:24.089989 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.002793550491333008 sec
W0605 17:37:24.093266 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004112720489501953 sec
W0605 17:37:24.094230 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003039836883544922 sec
W0605 17:37:24.095107 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00029397010803222656 sec
W0605 17:37:24.095797 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00030875205993652344 sec
W0605 17:37:24.143922 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033354759216308594 sec
W0605 17:37:24.145487 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003466606140136719 sec
W0605 17:37:24.157827 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036215782165527344 sec
W0605 17:37:24.227214 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00019693374633789062 sec
W0605 17:37:24.228162 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0016870498657226562 sec
W0605 17:37:24.229197 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.0030333995819091797 sec
W0605 17:37:24.232452 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.000308990478515625 sec
W0605 17:37:24.233449 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033164024353027344 sec
W0605 17:37:24.234340 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00029277801513671875 sec
W0605 17:37:24.235043 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003159046173095703 sec
W0605 17:37:24.341019 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.0011565685272216797 sec
W0605 17:37:24.451697 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00019693374633789062 sec
W0605 17:37:24.452555 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0015981197357177734 sec
W0605 17:37:24.453699 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.003048419952392578 sec
W0605 17:37:24.456939 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003116130828857422 sec
W0605 17:37:24.457936 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00031113624572753906 sec
W0605 17:37:24.458815 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002963542938232422 sec
W0605 17:37:24.459515 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003075599670410156 sec
W0605 17:37:24.566185 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_random_bits_original for pjit in 0.00116729736328125 sec
W0605 17:37:24.644937 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(init_fn) in 0.7951579093933105 sec
W0605 17:37:39.883617 139756486557120 dispatch.py:272] Finished XLA compilation of jit(init_fn) in 15.23807430267334 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005ac380 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005ac380, semaphore=0x6060001683e0, value=10 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=11, fence=0x604000ef0a10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005ac380, from_fence=0x606000195620 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000193460, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0007f6d70, f=0, wait_fence=0x6060005ac380 {0x6060001683e0:10}, signal_fence=0x604000ef0a10 {0x6060001683e0:11} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600079a8a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600079a900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a8a0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a900, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600079a960 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600079a9c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a960, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a9c0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600079aa20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600075fbc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079aa20, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600075fbc0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681620 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b00c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681620, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b00c0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006f03e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006811a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006f03e0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006811a0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002b7720 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681ec0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002b7720, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681ec0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000680de0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681e60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000680de0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681e60, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b0600 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000682940 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0600, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000682940, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b2040 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b25e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2040, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b25e0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681b00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b0540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681b00, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0540, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b2100 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006824c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2100, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006824c0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002b77e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b1c80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002b77e0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b1c80, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b1440 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b0180 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b1440, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0180, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b1860 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b1860, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681680, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000682ca0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b2ac0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000682ca0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2ac0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000681a40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b0f00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681a40, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0f00, semaphore=0x6060001683e0, value=11 (OK)
I0605 17:37:40.306165 139756486557120 trainer_lib.py:398] initial_vars: {'params': {'lm': {'final_ln': {'bias': (2048,), 'scale': (2048,)}, 'position_emb': {'emb_var': (2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (51200,)}, 'linear': {'w': (2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (24, 8192)}, 'linear': {'w': (24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (24, 2048)}, 'linear': {'w': (24, 8192, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}, 'self_attention': {'combined_qkv': {'w': (24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (24, 64)}, 'post': {'w': (24, 2048, 32, 64)}}}}}}}}}
W0605 17:37:40.307132 139756486557120 optimizers.py:1170] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
I0605 17:37:40.307184 139756486557120 optimizers.py:1173] Using sharded_adam.
W0605 17:37:40.307222 139756486557120 optimizers.py:580] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0002adb40 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:12}, signal={0x6060001682c0:13} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:13}, signal={0x6060001682c0:14} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bbe920 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600081a4c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bbe920, semaphore=0x6060001682c0, value=14 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600081a4c0, semaphore=0x6060001682c0, value=14 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0002b81c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:14}, signal={0x6060001682c0:15} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:15}, signal={0x6060001682c0:16} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bbea40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000753b00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bbea40, semaphore=0x6060001682c0, value=16 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000753b00, semaphore=0x6060001682c0, value=16 (OK)
W0605 17:37:40.309454 139756486557120 dispatch.py:272] Finished tracing + transforming jit(convert_element_type) in 0.0002548694610595703 sec
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0002b8700 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:16}, signal={0x6060001682c0:17} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:17}, signal={0x6060001682c0:18} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a9e080 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a9e1a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a9e080, semaphore=0x6060001682c0, value=18 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a9e1a0, semaphore=0x6060001682c0, value=18 (OK)
W0605 17:37:40.310383 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00022101402282714844 sec
W0605 17:37:40.310681 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:40.314482 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003664731979370117 sec
W0605 17:37:40.673562 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.35874080657958984 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003416c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003416c0, semaphore=0x6060001683e0, value=11 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=12, fence=0x604000f1d790 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003416c0, from_fence=0x606000a9e080 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a9e1a0, semaphore=0x6060001683e0, value=12 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x6060003416c0 {0x6060001683e0:11, 0x6060001682c0:18}, signal_fence=0x604000f1d790 {0x6060001683e0:12} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b29a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b2b20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b29a0, semaphore=0x6060001683e0, value=12 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2b20, semaphore=0x6060001683e0, value=12 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0002b87c0, wait={0x6060001682c0:18, 0x6060001683e0:12}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a55f80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:18}, signal={0x6060001682c0:19} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:19}, signal={0x6060001682c0:20} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b07200 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b07140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b07200, semaphore=0x6060001682c0, value=20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b07140, semaphore=0x6060001682c0, value=20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b070e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b070e0, semaphore=0x6060001683e0, value=12 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=13, fence=0x604000aa2c10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000b070e0, from_fence=0x606000b07200 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b07140, semaphore=0x6060001683e0, value=13 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000b070e0 {0x6060001683e0:12, 0x6060001682c0:20}, signal_fence=0x604000aa2c10 {0x6060001683e0:13} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005b9700 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005ba060 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005b9700, semaphore=0x6060001683e0, value=13 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005ba060, semaphore=0x6060001683e0, value=13 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000a55ec0, wait={0x6060001682c0:20, 0x6060001683e0:13}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a28f80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:20}, signal={0x6060001682c0:21} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:21}, signal={0x6060001682c0:22} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000972b00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009712a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000972b00, semaphore=0x6060001682c0, value=22 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009712a0, semaphore=0x6060001682c0, value=22 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600096ffe0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600096ffe0, semaphore=0x6060001683e0, value=13 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=14, fence=0x6040008e9b10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600096ffe0, from_fence=0x606000972b00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009712a0, semaphore=0x6060001683e0, value=14 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x60600096ffe0 {0x6060001683e0:13, 0x6060001682c0:22}, signal_fence=0x6040008e9b10 {0x6060001683e0:14} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000970700 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600034d420 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000970700, semaphore=0x6060001683e0, value=14 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600034d420, semaphore=0x6060001683e0, value=14 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000a29040, wait={0x6060001682c0:22, 0x6060001683e0:14}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a294c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:22}, signal={0x6060001682c0:23} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:23}, signal={0x6060001682c0:24} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000cebe80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000cebe20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cebe80, semaphore=0x6060001682c0, value=24 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cebe20, semaphore=0x6060001682c0, value=24 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000946460 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000946460, semaphore=0x6060001683e0, value=14 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=15, fence=0x604001474b90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000946460, from_fence=0x606000cebe80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cebe20, semaphore=0x6060001683e0, value=15 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000946460 {0x6060001683e0:14, 0x6060001682c0:24}, signal_fence=0x604001474b90 {0x6060001683e0:15} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000946700 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009467c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000946700, semaphore=0x6060001683e0, value=15 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009467c0, semaphore=0x6060001683e0, value=15 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000a29580, wait={0x6060001682c0:24, 0x6060001683e0:15}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a28080 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:24}, signal={0x6060001682c0:25} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:25}, signal={0x6060001682c0:26} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009472a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000661b20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009472a0, semaphore=0x6060001682c0, value=26 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000661b20, semaphore=0x6060001682c0, value=26 (OK)
W0605 17:37:40.679012 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00041937828063964844 sec
W0605 17:37:40.679480 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:40.684844 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0051648616790771484 sec
W0605 17:37:41.056386 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3710479736328125 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000967100 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000967100, semaphore=0x6060001683e0, value=15 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=16, fence=0x60400162c990 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000967100, from_fence=0x6060009472a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000661b20, semaphore=0x6060001683e0, value=16 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000469ec0, f=0, wait_fence=0x606000967100 {0x6060001683e0:15, 0x6060001682c0:26}, signal_fence=0x60400162c990 {0x6060001683e0:16} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000967e80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000967e20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000967e80, semaphore=0x6060001683e0, value=16 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000967e20, semaphore=0x6060001683e0, value=16 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000a28140, wait={0x6060001682c0:26, 0x6060001683e0:16}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0005de080 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:26}, signal={0x6060001682c0:27} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:27}, signal={0x6060001682c0:28} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000968cc0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000968c60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000968cc0, semaphore=0x6060001682c0, value=28 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000968c60, semaphore=0x6060001682c0, value=28 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007e8300 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007e8300, semaphore=0x6060001683e0, value=16 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=17, fence=0x60400117b350 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060007e8300, from_fence=0x606000968cc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000968c60, semaphore=0x6060001683e0, value=17 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000469ec0, f=0, wait_fence=0x6060007e8300 {0x6060001683e0:16, 0x6060001682c0:28}, signal_fence=0x60400117b350 {0x6060001683e0:17} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007e8240 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000761a80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007e8240, semaphore=0x6060001683e0, value=17 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000761a80, semaphore=0x6060001683e0, value=17 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b9d400, wait={0x6060001682c0:28, 0x6060001683e0:17}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000471e80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:28}, signal={0x6060001682c0:29} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:29}, signal={0x6060001682c0:30} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600084b120 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600084b0c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600084b120, semaphore=0x6060001682c0, value=30 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600084b0c0, semaphore=0x6060001682c0, value=30 (OK)
W0605 17:37:41.059996 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002830028533935547 sec
W0605 17:37:41.060308 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:41.064166 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003727436065673828 sec
W0605 17:37:41.432264 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.36777281761169434 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c9b000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c9b000, semaphore=0x6060001683e0, value=17 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=18, fence=0x6040016a0110 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c9b000, from_fence=0x60600084b120 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600084b0c0, semaphore=0x6060001683e0, value=18 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0008a7ea0, f=0, wait_fence=0x606000c9b000 {0x6060001683e0:17, 0x6060001682c0:30}, signal_fence=0x6040016a0110 {0x6060001683e0:18} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ca0c40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006442a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ca0c40, semaphore=0x6060001683e0, value=18 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006442a0, semaphore=0x6060001683e0, value=18 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b61e80, wait={0x6060001682c0:30, 0x6060001683e0:18}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000266200 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:30}, signal={0x6060001682c0:31} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:31}, signal={0x6060001682c0:32} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a7a20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000391dc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a7a20, semaphore=0x6060001682c0, value=32 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000391dc0, semaphore=0x6060001682c0, value=32 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003902c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003902c0, semaphore=0x6060001683e0, value=18 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=19, fence=0x60400041ba10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003902c0, from_fence=0x6060003a7a20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000391dc0, semaphore=0x6060001683e0, value=19 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0008a7ea0, f=0, wait_fence=0x6060003902c0 {0x6060001683e0:18, 0x6060001682c0:32}, signal_fence=0x60400041ba10 {0x6060001683e0:19} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003915e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001c5e60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003915e0, semaphore=0x6060001683e0, value=19 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001c5e60, semaphore=0x6060001683e0, value=19 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000352240, wait={0x6060001682c0:32, 0x6060001683e0:19}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a98e80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:32}, signal={0x6060001682c0:33} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:33}, signal={0x6060001682c0:34} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000145880 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bf8ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000145880, semaphore=0x6060001682c0, value=34 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bf8ee0, semaphore=0x6060001682c0, value=34 (OK)
W0605 17:37:41.435681 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0003001689910888672 sec
W0605 17:37:41.435994 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:41.439823 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003693819046020508 sec
W0605 17:37:41.806956 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.36679840087890625 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000981c80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000981c80, semaphore=0x6060001683e0, value=19 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=20, fence=0x6040014c8650 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000981c80, from_fence=0x606000145880 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bf8ee0, semaphore=0x6060001683e0, value=20 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0005eb9c0, f=0, wait_fence=0x606000981c80 {0x6060001683e0:19, 0x6060001682c0:34}, signal_fence=0x6040014c8650 {0x6060001683e0:20} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005719a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007cf6a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005719a0, semaphore=0x6060001683e0, value=20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007cf6a0, semaphore=0x6060001683e0, value=20 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000676c00, wait={0x6060001682c0:34, 0x6060001683e0:20}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000443080 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:34}, signal={0x6060001682c0:35} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:35}, signal={0x6060001682c0:36} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005da0a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000456140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005da0a0, semaphore=0x6060001682c0, value=36 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000456140, semaphore=0x6060001682c0, value=36 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600082fb20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600082fb20, semaphore=0x6060001683e0, value=20 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=21, fence=0x60400161db50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600082fb20, from_fence=0x6060005da0a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000456140, semaphore=0x6060001683e0, value=21 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0005eb9c0, f=0, wait_fence=0x60600082fb20 {0x6060001683e0:20, 0x6060001682c0:36}, signal_fence=0x60400161db50 {0x6060001683e0:21} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600044a9e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600061e4a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600044a9e0, semaphore=0x6060001683e0, value=21 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600061e4a0, semaphore=0x6060001683e0, value=21 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000ad9ec0, wait={0x6060001682c0:36, 0x6060001683e0:21}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0009f4300 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:36}, signal={0x6060001682c0:37} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:37}, signal={0x6060001682c0:38} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a93c40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a93820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a93c40, semaphore=0x6060001682c0, value=38 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a93820, semaphore=0x6060001682c0, value=38 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0002c6ec0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:38}, signal={0x6060001682c0:39} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:39}, signal={0x6060001682c0:40} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000927e00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000315740 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000927e00, semaphore=0x6060001682c0, value=40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000315740, semaphore=0x6060001682c0, value=40 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000c33280 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:40}, signal={0x6060001682c0:41} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:41}, signal={0x6060001682c0:42} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ba93e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007beea0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ba93e0, semaphore=0x6060001682c0, value=42 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007beea0, semaphore=0x6060001682c0, value=42 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00039c700 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:42}, signal={0x6060001682c0:43} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:43}, signal={0x6060001682c0:44} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b2c220 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c427a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b2c220, semaphore=0x6060001682c0, value=44 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c427a0, semaphore=0x6060001682c0, value=44 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0005fba80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:44}, signal={0x6060001682c0:45} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:45}, signal={0x6060001682c0:46} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c42b00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c422c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c42b00, semaphore=0x6060001682c0, value=46 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c422c0, semaphore=0x6060001682c0, value=46 (OK)
W0605 17:37:41.817803 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00026488304138183594 sec
W0605 17:37:41.818148 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:41.822030 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0037508010864257812 sec
W0605 17:37:42.176656 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3542964458465576 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8880 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8880, semaphore=0x6060001683e0, value=21 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=22, fence=0x60400096c810 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060008c8880, from_fence=0x606000c42b00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c422c0, semaphore=0x6060001683e0, value=22 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a12c90, f=0, wait_fence=0x6060008c8880 {0x6060001683e0:21, 0x6060001682c0:46}, signal_fence=0x60400096c810 {0x6060001683e0:22} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003297e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c87c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003297e0, semaphore=0x6060001683e0, value=22 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c87c0, semaphore=0x6060001683e0, value=22 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0005fb600, wait={0x6060001682c0:46, 0x6060001683e0:22}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000189c40 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:46}, signal={0x6060001682c0:47} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:47}, signal={0x6060001682c0:48} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8f40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8c40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8f40, semaphore=0x6060001682c0, value=48 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8c40, semaphore=0x6060001682c0, value=48 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8fa0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8fa0, semaphore=0x6060001683e0, value=22 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=23, fence=0x604000bb7e10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060008c8fa0, from_fence=0x6060008c8f40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8c40, semaphore=0x6060001683e0, value=23 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a12c90, f=0, wait_fence=0x6060008c8fa0 {0x6060001683e0:22, 0x6060001682c0:48}, signal_fence=0x604000bb7e10 {0x6060001683e0:23} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8700 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8dc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8700, semaphore=0x6060001683e0, value=23 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8dc0, semaphore=0x6060001683e0, value=23 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0006394c0, wait={0x6060001682c0:48, 0x6060001683e0:23}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b48500 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:48}, signal={0x6060001682c0:49} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:49}, signal={0x6060001682c0:50} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c7f80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c8040 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c7f80, semaphore=0x6060001682c0, value=50 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8040, semaphore=0x6060001682c0, value=50 (OK)
W0605 17:37:42.180253 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00028324127197265625 sec
W0605 17:37:42.180571 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:42.184416 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003713846206665039 sec
W0605 17:37:42.550396 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.365649938583374 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000cf2000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cf2000, semaphore=0x6060001683e0, value=23 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=24, fence=0x6040002a0190 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000cf2000, from_fence=0x6060008c7f80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8040, semaphore=0x6060001683e0, value=24 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000710b30, f=0, wait_fence=0x606000cf2000 {0x6060001683e0:23, 0x6060001682c0:50}, signal_fence=0x6040002a0190 {0x6060001683e0:24} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600095efa0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000342920 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600095efa0, semaphore=0x6060001683e0, value=24 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000342920, semaphore=0x6060001683e0, value=24 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0006397c0, wait={0x6060001682c0:50, 0x6060001683e0:24}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00052ef40 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:50}, signal={0x6060001682c0:51} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:51}, signal={0x6060001682c0:52} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000821540 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009fe960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000821540, semaphore=0x6060001682c0, value=52 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009fe960, semaphore=0x6060001682c0, value=52 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d6e320 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d6e320, semaphore=0x6060001683e0, value=24 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=25, fence=0x6040009082d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000d6e320, from_fence=0x606000821540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009fe960, semaphore=0x6060001683e0, value=25 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000710b30, f=0, wait_fence=0x606000d6e320 {0x6060001683e0:24, 0x6060001682c0:52}, signal_fence=0x6040009082d0 {0x6060001683e0:25} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c5820 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d6e5c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c5820, semaphore=0x6060001683e0, value=25 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d6e5c0, semaphore=0x6060001683e0, value=25 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00052f000, wait={0x6060001682c0:52, 0x6060001683e0:25}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00052e880 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:52}, signal={0x6060001682c0:53} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:53}, signal={0x6060001682c0:54} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008d3320 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a3da20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008d3320, semaphore=0x6060001682c0, value=54 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a3da20, semaphore=0x6060001682c0, value=54 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000131c00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000131c00, semaphore=0x6060001683e0, value=25 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=26, fence=0x6040006bcc10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000131c00, from_fence=0x6060008d3320 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a3da20, semaphore=0x6060001683e0, value=26 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000131c00 {0x6060001683e0:25, 0x6060001682c0:54}, signal_fence=0x6040006bcc10 {0x6060001683e0:26} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c2da40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009fe600 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c2da40, semaphore=0x6060001683e0, value=26 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009fe600, semaphore=0x6060001683e0, value=26 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00052e940, wait={0x6060001682c0:54, 0x6060001683e0:26}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000bc80c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:54}, signal={0x6060001682c0:55} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:55}, signal={0x6060001682c0:56} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d6db40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009823a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d6db40, semaphore=0x6060001682c0, value=56 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009823a0, semaphore=0x6060001682c0, value=56 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000732260 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000732260, semaphore=0x6060001683e0, value=26 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=27, fence=0x604000248d10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000732260, from_fence=0x606000d6db40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009823a0, semaphore=0x6060001683e0, value=27 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000732260 {0x6060001683e0:26, 0x6060001682c0:56}, signal_fence=0x604000248d10 {0x6060001683e0:27} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d4b580 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000254d20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d4b580, semaphore=0x6060001683e0, value=27 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000254d20, semaphore=0x6060001683e0, value=27 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000bc4880, wait={0x6060001682c0:56, 0x6060001683e0:27}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000bc6800 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:56}, signal={0x6060001682c0:57} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:57}, signal={0x6060001682c0:58} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b7e240 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b01b60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b7e240, semaphore=0x6060001682c0, value=58 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b01b60, semaphore=0x6060001682c0, value=58 (OK)
W0605 17:37:42.556579 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002880096435546875 sec
W0605 17:37:42.556906 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:42.560795 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003753662109375 sec
W0605 17:37:42.913717 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.35259294509887695 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600032d3e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600032d3e0, semaphore=0x6060001683e0, value=27 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=28, fence=0x60400054ed50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600032d3e0, from_fence=0x606000b7e240 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b01b60, semaphore=0x6060001683e0, value=28 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a0ef60, f=0, wait_fence=0x60600032d3e0 {0x6060001683e0:27, 0x6060001682c0:58}, signal_fence=0x60400054ed50 {0x6060001683e0:28} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072e9c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037f940 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072e9c0, semaphore=0x6060001683e0, value=28 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f940, semaphore=0x6060001683e0, value=28 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000bc5000, wait={0x6060001682c0:58, 0x6060001683e0:28}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b8b580 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:58}, signal={0x6060001682c0:59} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:59}, signal={0x6060001682c0:60} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600036e900 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037f700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600036e900, semaphore=0x6060001682c0, value=60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f700, semaphore=0x6060001682c0, value=60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072e540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072e540, semaphore=0x6060001683e0, value=28 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=29, fence=0x604001347e50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600072e540, from_fence=0x60600036e900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f700, semaphore=0x6060001683e0, value=29 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a0ef60, f=0, wait_fence=0x60600072e540 {0x6060001683e0:28, 0x6060001682c0:60}, signal_fence=0x604001347e50 {0x6060001683e0:29} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600036e060 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072eae0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600036e060, semaphore=0x6060001683e0, value=29 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072eae0, semaphore=0x6060001683e0, value=29 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b8b1c0, wait={0x6060001682c0:60, 0x6060001683e0:29}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0005eb1c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:60}, signal={0x6060001682c0:61} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:61}, signal={0x6060001682c0:62} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037f640 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037f820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f640, semaphore=0x6060001682c0, value=62 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f820, semaphore=0x6060001682c0, value=62 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000006680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000006680, semaphore=0x6060001683e0, value=29 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=30, fence=0x604001319f90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000006680, from_fence=0x60600037f640 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f820, semaphore=0x6060001683e0, value=30 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000006680 {0x6060001683e0:29, 0x6060001682c0:62}, signal_fence=0x604001319f90 {0x6060001683e0:30} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003295a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000329660 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003295a0, semaphore=0x6060001683e0, value=30 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000329660, semaphore=0x6060001683e0, value=30 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0005eb7c0, wait={0x6060001682c0:62, 0x6060001683e0:30}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0002c7340 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:62}, signal={0x6060001682c0:63} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:63}, signal={0x6060001682c0:64} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037d720 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000129a40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037d720, semaphore=0x6060001682c0, value=64 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000129a40, semaphore=0x6060001682c0, value=64 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000129980 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000129980, semaphore=0x6060001683e0, value=30 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=31, fence=0x60400131fb90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000129980, from_fence=0x60600037d720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000129a40, semaphore=0x6060001683e0, value=31 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000129980 {0x6060001683e0:30, 0x6060001682c0:64}, signal_fence=0x60400131fb90 {0x6060001683e0:31} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000129680 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001cd8a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000129680, semaphore=0x6060001683e0, value=31 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001cd8a0, semaphore=0x6060001683e0, value=31 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0002c7400, wait={0x6060001682c0:64, 0x6060001683e0:31}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000054280 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:64}, signal={0x6060001682c0:65} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:65}, signal={0x6060001682c0:66} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000cf020 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000c76a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000cf020, semaphore=0x6060001682c0, value=66 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000c76a0, semaphore=0x6060001682c0, value=66 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037d480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037d480, semaphore=0x6060001683e0, value=31 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=32, fence=0x604000ab0c10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600037d480, from_fence=0x6060000cf020 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000c76a0, semaphore=0x6060001683e0, value=32 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x60600037d480 {0x6060001683e0:31, 0x6060001682c0:66}, signal_fence=0x604000ab0c10 {0x6060001683e0:32} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000006860 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001a9120 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000006860, semaphore=0x6060001683e0, value=32 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001a9120, semaphore=0x6060001683e0, value=32 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000332980, wait={0x6060001682c0:66, 0x6060001683e0:32}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000ba52c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:66}, signal={0x6060001682c0:67} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:67}, signal={0x6060001682c0:68} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000192440 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600075ffe0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000192440, semaphore=0x6060001682c0, value=68 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600075ffe0, semaphore=0x6060001682c0, value=68 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600052df60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600052df60, semaphore=0x6060001683e0, value=32 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=33, fence=0x604000aae850 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600052df60, from_fence=0x606000192440 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600075ffe0, semaphore=0x6060001683e0, value=33 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x60600052df60 {0x6060001683e0:32, 0x6060001682c0:68}, signal_fence=0x604000aae850 {0x6060001683e0:33} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000a1780 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600052e1a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a1780, semaphore=0x6060001683e0, value=33 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600052e1a0, semaphore=0x6060001683e0, value=33 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000344200, wait={0x6060001682c0:68, 0x6060001683e0:33}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0006832c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:68}, signal={0x6060001682c0:69} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:69}, signal={0x6060001682c0:70} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccb480 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccb540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb480, semaphore=0x6060001682c0, value=70 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb540, semaphore=0x6060001682c0, value=70 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccb3c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb3c0, semaphore=0x6060001683e0, value=33 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=34, fence=0x604000dd5110 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000ccb3c0, from_fence=0x606000ccb480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb540, semaphore=0x6060001683e0, value=34 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000ccb3c0 {0x6060001683e0:33, 0x6060001682c0:70}, signal_fence=0x604000dd5110 {0x6060001683e0:34} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccbb40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccb720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbb40, semaphore=0x6060001683e0, value=34 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb720, semaphore=0x6060001683e0, value=34 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00022d440, wait={0x6060001682c0:70, 0x6060001683e0:34}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000a85800 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:70}, signal={0x6060001682c0:71} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:71}, signal={0x6060001682c0:72} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccbf60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccbe40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbf60, semaphore=0x6060001682c0, value=72 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbe40, semaphore=0x6060001682c0, value=72 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccbd80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbd80, semaphore=0x6060001683e0, value=34 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=35, fence=0x604000727d10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000ccbd80, from_fence=0x606000ccbf60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbe40, semaphore=0x6060001683e0, value=35 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x606000ccbd80 {0x6060001683e0:34, 0x6060001682c0:72}, signal_fence=0x604000727d10 {0x6060001683e0:35} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccbcc0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072c140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccbcc0, semaphore=0x6060001683e0, value=35 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072c140, semaphore=0x6060001683e0, value=35 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000079e40, wait={0x6060001682c0:72, 0x6060001683e0:35}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00006b5c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:72}, signal={0x6060001682c0:73} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:73}, signal={0x6060001682c0:74} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004d3c60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600025a0c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004d3c60, semaphore=0x6060001682c0, value=74 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600025a0c0, semaphore=0x6060001682c0, value=74 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006fe8a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006fe8a0, semaphore=0x6060001683e0, value=35 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=36, fence=0x6040007284d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006fe8a0, from_fence=0x6060004d3c60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600025a0c0, semaphore=0x6060001683e0, value=36 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x6060006fe8a0 {0x6060001683e0:35, 0x6060001682c0:74}, signal_fence=0x6040007284d0 {0x6060001683e0:36} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072c080 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005f5ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072c080, semaphore=0x6060001683e0, value=36 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005f5ee0, semaphore=0x6060001683e0, value=36 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0004a8780, wait={0x6060001682c0:74, 0x6060001683e0:36}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000676900 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:74}, signal={0x6060001682c0:75} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:75}, signal={0x6060001682c0:76} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004d30c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007d5520 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004d30c0, semaphore=0x6060001682c0, value=76 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007d5520, semaphore=0x6060001682c0, value=76 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003996e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003996e0, semaphore=0x6060001683e0, value=36 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=37, fence=0x604000231e10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003996e0, from_fence=0x6060004d30c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007d5520, semaphore=0x6060001683e0, value=37 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004d9af0, f=0, wait_fence=0x6060003996e0 {0x6060001683e0:36, 0x6060001682c0:76}, signal_fence=0x604000231e10 {0x6060001683e0:37} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600036fbc0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008daca0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600036fbc0, semaphore=0x6060001683e0, value=37 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008daca0, semaphore=0x6060001683e0, value=37 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009b5cc0, wait={0x6060001682c0:76, 0x6060001683e0:37}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000685000 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:76}, signal={0x6060001682c0:77} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:77}, signal={0x6060001682c0:78} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000978da0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000978440 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000978da0, semaphore=0x6060001682c0, value=78 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000978440, semaphore=0x6060001682c0, value=78 (OK)
W0605 17:37:42.924442 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002827644348144531 sec
W0605 17:37:42.924763 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:42.928637 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0037441253662109375 sec
W0605 17:37:43.298731 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.36975932121276855 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c3b300 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c3b300, semaphore=0x6060001683e0, value=37 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=38, fence=0x6040015e6290 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c3b300, from_fence=0x606000978da0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000978440, semaphore=0x6060001683e0, value=38 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0008a2490, f=0, wait_fence=0x606000c3b300 {0x6060001683e0:37, 0x6060001682c0:78}, signal_fence=0x6040015e6290 {0x6060001683e0:38} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c44960 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600018fb60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c44960, semaphore=0x6060001683e0, value=38 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600018fb60, semaphore=0x6060001683e0, value=38 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000bd0f40, wait={0x6060001682c0:78, 0x6060001683e0:38}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0009b2a80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:78}, signal={0x6060001682c0:79} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:79}, signal={0x6060001682c0:80} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000628be0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b6f540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000628be0, semaphore=0x6060001682c0, value=80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b6f540, semaphore=0x6060001682c0, value=80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000628e80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000628e80, semaphore=0x6060001683e0, value=38 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=39, fence=0x604000159250 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000628e80, from_fence=0x606000628be0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b6f540, semaphore=0x6060001683e0, value=39 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0008a2490, f=0, wait_fence=0x606000628e80 {0x6060001683e0:38, 0x6060001682c0:80}, signal_fence=0x604000159250 {0x6060001683e0:39} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b37da0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b5400 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b37da0, semaphore=0x6060001683e0, value=39 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b5400, semaphore=0x6060001683e0, value=39 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009b2b40, wait={0x6060001682c0:80, 0x6060001683e0:39}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0009b2fc0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:80}, signal={0x6060001682c0:81} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:81}, signal={0x6060001682c0:82} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d64480 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d64360 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d64480, semaphore=0x6060001682c0, value=82 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d64360, semaphore=0x6060001682c0, value=82 (OK)
W0605 17:37:43.303542 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002834796905517578 sec
W0605 17:37:43.303862 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:43.307848 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0038557052612304688 sec
W0605 17:37:43.702368 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3941676616668701 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006c8960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006c8960, semaphore=0x6060001683e0, value=39 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=40, fence=0x6040015f1f90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006c8960, from_fence=0x606000d64480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d64360, semaphore=0x6060001683e0, value=40 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00035ca30, f=0, wait_fence=0x6060006c8960 {0x6060001683e0:39, 0x6060001682c0:82}, signal_fence=0x6040015f1f90 {0x6060001683e0:40} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bb9ee0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bb9fa0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bb9ee0, semaphore=0x6060001683e0, value=40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bb9fa0, semaphore=0x6060001683e0, value=40 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00098f080, wait={0x6060001682c0:82, 0x6060001683e0:40}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b13100 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:82}, signal={0x6060001682c0:83} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:83}, signal={0x6060001682c0:84} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bb93a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d0f820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bb93a0, semaphore=0x6060001682c0, value=84 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d0f820, semaphore=0x6060001682c0, value=84 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003d5500 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003d5500, semaphore=0x6060001683e0, value=40 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=41, fence=0x604001406190 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003d5500, from_fence=0x606000bb93a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d0f820, semaphore=0x6060001683e0, value=41 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00035ca30, f=0, wait_fence=0x6060003d5500 {0x6060001683e0:40, 0x6060001682c0:84}, signal_fence=0x604001406190 {0x6060001683e0:41} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000283ac0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003d5aa0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000283ac0, semaphore=0x6060001683e0, value=41 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003d5aa0, semaphore=0x6060001683e0, value=41 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000768340, wait={0x6060001682c0:84, 0x6060001683e0:41}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000c4a2c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:84}, signal={0x6060001682c0:85} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:85}, signal={0x6060001682c0:86} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600050d080 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600050d0e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600050d080, semaphore=0x6060001682c0, value=86 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600050d0e0, semaphore=0x6060001682c0, value=86 (OK)
W0605 17:37:43.708622 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0005459785461425781 sec
W0605 17:37:43.709253 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:43.713417 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0039212703704833984 sec
W0605 17:37:44.077934 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3641831874847412 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002b71e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002b71e0, semaphore=0x6060001683e0, value=41 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=42, fence=0x604000450490 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060002b71e0, from_fence=0x60600050d080 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600050d0e0, semaphore=0x6060001683e0, value=42 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0007efaa0, f=0, wait_fence=0x6060002b71e0 {0x6060001683e0:41, 0x6060001682c0:86}, signal_fence=0x604000450490 {0x6060001683e0:42} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000e9240 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005bfb20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000e9240, semaphore=0x6060001683e0, value=42 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005bfb20, semaphore=0x6060001683e0, value=42 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b7c1c0, wait={0x6060001682c0:86, 0x6060001683e0:42}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b42c80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:86}, signal={0x6060001682c0:87} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:87}, signal={0x6060001682c0:88} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005f4e00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600026dce0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005f4e00, semaphore=0x6060001682c0, value=88 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600026dce0, semaphore=0x6060001682c0, value=88 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005821a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005821a0, semaphore=0x6060001683e0, value=42 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=43, fence=0x6040011e9d10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005821a0, from_fence=0x6060005f4e00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600026dce0, semaphore=0x6060001683e0, value=43 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0007efaa0, f=0, wait_fence=0x6060005821a0 {0x6060001683e0:42, 0x6060001682c0:88}, signal_fence=0x6040011e9d10 {0x6060001683e0:43} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005b2e60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600073d120 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005b2e60, semaphore=0x6060001683e0, value=43 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600073d120, semaphore=0x6060001683e0, value=43 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0006556c0, wait={0x6060001682c0:88, 0x6060001683e0:43}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000655600 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:88}, signal={0x6060001682c0:89} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:89}, signal={0x6060001682c0:90} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600029f840 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000609f20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600029f840, semaphore=0x6060001682c0, value=90 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000609f20, semaphore=0x6060001682c0, value=90 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00006ce80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:90}, signal={0x6060001682c0:91} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:91}, signal={0x6060001682c0:92} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bccc60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c7aa20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bccc60, semaphore=0x6060001682c0, value=92 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c7aa20, semaphore=0x6060001682c0, value=92 (OK)
W0605 17:37:44.082766 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002627372741699219 sec
W0605 17:37:44.083084 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(int32[])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:44.087130 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003915309906005859 sec
W0605 17:37:44.459510 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3720529079437256 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009638c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009638c0, semaphore=0x6060001683e0, value=43 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=44, fence=0x604000ddb850 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060009638c0, from_fence=0x606000ba93e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007beea0, semaphore=0x6060001683e0, value=44 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00011f9b0, f=0, wait_fence=0x6060009638c0 {0x6060001683e0:43, 0x6060001682c0:42}, signal_fence=0x604000ddb850 {0x6060001683e0:44} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000965a80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006dbb00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000965a80, semaphore=0x6060001683e0, value=44 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006dbb00, semaphore=0x6060001683e0, value=44 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006dbe00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006dbe00, semaphore=0x6060001683e0, value=44 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=45, fence=0x604000b1dc50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006dbe00, from_fence=0x606000b2c220 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c427a0, semaphore=0x6060001683e0, value=45 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00011f9b0, f=0, wait_fence=0x6060006dbe00 {0x6060001683e0:44, 0x6060001682c0:44}, signal_fence=0x604000b1dc50 {0x6060001683e0:45} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008210c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600086c7e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008210c0, semaphore=0x6060001683e0, value=45 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600086c7e0, semaphore=0x6060001683e0, value=45 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600015abe0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600015abe0, semaphore=0x6060001683e0, value=45 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=46, fence=0x60400013e310 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600015abe0, from_fence=0x60600029f840 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000609f20, semaphore=0x6060001683e0, value=46 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00011f9b0, f=0, wait_fence=0x60600015abe0 {0x6060001683e0:45, 0x6060001682c0:90}, signal_fence=0x60400013e310 {0x6060001683e0:46} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007a8160 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000503c60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007a8160, semaphore=0x6060001683e0, value=46 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000503c60, semaphore=0x6060001683e0, value=46 (OK)
W0605 17:37:44.462649 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0002999305725097656 sec
W0605 17:37:44.462978 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[8192])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:44.466973 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.0038557052612304688 sec
W0605 17:37:44.827181 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.359877347946167 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000981e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000981e0, semaphore=0x6060001683e0, value=46 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=47, fence=0x604000d97d10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060000981e0, from_fence=0x6060003297e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c87c0, semaphore=0x6060001683e0, value=47 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0007d79a0, f=0, wait_fence=0x6060000981e0 {0x6060001683e0:46}, signal_fence=0x604000d97d10 {0x6060001683e0:47} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600009c7a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000779d80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600009c7a0, semaphore=0x6060001683e0, value=47 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000779d80, semaphore=0x6060001683e0, value=47 (OK)
W0605 17:37:44.829687 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0003190040588378906 sec
W0605 17:37:44.830013 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[2048,8192])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:44.834094 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.003939390182495117 sec
W0605 17:37:45.209967 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.37554359436035156 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000356660 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000356660, semaphore=0x6060001683e0, value=47 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=48, fence=0x604000ddf9d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000356660, from_fence=0x60600095efa0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000342920, semaphore=0x6060001683e0, value=48 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0003d73a0, f=0, wait_fence=0x606000356660 {0x6060001683e0:47}, signal_fence=0x604000ddf9d0 {0x6060001683e0:48} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004af600 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003eb820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004af600, semaphore=0x6060001683e0, value=48 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003eb820, semaphore=0x6060001683e0, value=48 (OK)
W0605 17:37:45.224115 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00031113624572753906 sec
W0605 17:37:45.224486 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[2048])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:45.228661 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.004036903381347656 sec
W0605 17:37:45.607018 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.37800025939941406 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005cca20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005cca20, semaphore=0x6060001683e0, value=48 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=49, fence=0x60400124fe10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005cca20, from_fence=0x606000c2da40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009fe600, semaphore=0x6060001683e0, value=49 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x6060005cca20 {0x6060001683e0:48}, signal_fence=0x60400124fe10 {0x6060001683e0:49} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600049fc40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004e61a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600049fc40, semaphore=0x6060001683e0, value=49 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004e61a0, semaphore=0x6060001683e0, value=49 (OK)
W0605 17:37:45.610910 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0005667209625244141 sec
W0605 17:37:45.611510 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[8192,2048])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:45.616619 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.00486302375793457 sec
W0605 17:37:45.954178 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.33722805976867676 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009c0920 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009c0920, semaphore=0x6060001683e0, value=49 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=50, fence=0x604000148c10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060009c0920, from_fence=0x60600072e9c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037f940, semaphore=0x6060001683e0, value=50 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000843b60, f=0, wait_fence=0x6060009c0920 {0x6060001683e0:49}, signal_fence=0x604000148c10 {0x6060001683e0:50} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004fbaa0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a06a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004fbaa0, semaphore=0x6060001683e0, value=50 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a06a0, semaphore=0x6060001683e0, value=50 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000844f40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000844f40, semaphore=0x6060001683e0, value=50 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=51, fence=0x6040008431d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000844f40, from_fence=0x6060003295a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000329660, semaphore=0x6060001683e0, value=51 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x606000844f40 {0x6060001683e0:50}, signal_fence=0x6040008431d0 {0x6060001683e0:51} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003f5780 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600089cfc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003f5780, semaphore=0x6060001683e0, value=51 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600089cfc0, semaphore=0x6060001683e0, value=51 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c407c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c407c0, semaphore=0x6060001683e0, value=51 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=52, fence=0x6040000f5a90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c407c0, from_fence=0x606000006860 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001a9120, semaphore=0x6060001683e0, value=52 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x606000c407c0 {0x6060001683e0:51}, signal_fence=0x6040000f5a90 {0x6060001683e0:52} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009c0200 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c41900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009c0200, semaphore=0x6060001683e0, value=52 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c41900, semaphore=0x6060001683e0, value=52 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c41960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c41960, semaphore=0x6060001683e0, value=52 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=53, fence=0x60400165c3d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c41960, from_fence=0x606000ccbb40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccb720, semaphore=0x6060001683e0, value=53 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x606000c41960 {0x6060001683e0:52}, signal_fence=0x60400165c3d0 {0x6060001683e0:53} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009c08c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009e2940 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009c08c0, semaphore=0x6060001683e0, value=53 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009e2940, semaphore=0x6060001683e0, value=53 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bc9d20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bc9d20, semaphore=0x6060001683e0, value=53 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=54, fence=0x60400165c5d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000bc9d20, from_fence=0x60600072c080 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005f5ee0, semaphore=0x6060001683e0, value=54 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x606000bc9d20 {0x6060001683e0:53}, signal_fence=0x60400165c5d0 {0x6060001683e0:54} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bca3e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bca5c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bca3e0, semaphore=0x6060001683e0, value=54 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bca5c0, semaphore=0x6060001683e0, value=54 (OK)
W0605 17:37:45.970173 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00032019615173339844 sec
W0605 17:37:45.970548 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[3,2048,32,64])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:45.974761 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.004075050354003906 sec
W0605 17:37:46.345563 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.370466947555542 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006592a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006592a0, semaphore=0x6060001683e0, value=54 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=55, fence=0x6040005d00d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006592a0, from_fence=0x606000c44960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600018fb60, semaphore=0x6060001683e0, value=55 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a36100, f=0, wait_fence=0x6060006592a0 {0x6060001683e0:54}, signal_fence=0x6040005d00d0 {0x6060001683e0:55} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600065b940 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000659900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600065b940, semaphore=0x6060001683e0, value=55 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000659900, semaphore=0x6060001683e0, value=55 (OK)
W0605 17:37:46.357086 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.00030922889709472656 sec
W0605 17:37:46.357457 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[64])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:46.361810 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.004218578338623047 sec
W0605 17:37:46.735365 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.3732173442840576 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007ea280 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007ea280, semaphore=0x6060001683e0, value=55 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=56, fence=0x6040014db3d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060007ea280, from_fence=0x606000bb9ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bb9fa0, semaphore=0x6060001683e0, value=56 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00036c520, f=0, wait_fence=0x6060007ea280 {0x6060001683e0:55}, signal_fence=0x6040014db3d0 {0x6060001683e0:56} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600023cde0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009040a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600023cde0, semaphore=0x6060001683e0, value=56 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009040a0, semaphore=0x6060001683e0, value=56 (OK)
W0605 17:37:46.737891 139756486557120 dispatch.py:272] Finished tracing + transforming jit(broadcast_in_dim) in 0.0003180503845214844 sec
W0605 17:37:46.738218 139756486557120 pxla.py:1882] Compiling broadcast_in_dim for with global shapes and types [ShapedArray(float32[2048,32,64])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:46.742366 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(broadcast_in_dim) in 0.004011631011962891 sec
W0605 17:37:47.107496 139756486557120 dispatch.py:272] Finished XLA compilation of jit(broadcast_in_dim) in 0.36479902267456055 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a78d00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a78d00, semaphore=0x6060001683e0, value=56 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=57, fence=0x6040011bac90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000a78d00, from_fence=0x6060000e9240 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005bfb20, semaphore=0x6060001683e0, value=57 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0003c8250, f=0, wait_fence=0x606000a78d00 {0x6060001683e0:56}, signal_fence=0x6040011bac90 {0x6060001683e0:57} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001619c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002a2c00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001619c0, semaphore=0x6060001683e0, value=57 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002a2c00, semaphore=0x6060001683e0, value=57 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600086a9e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600086a9e0, semaphore=0x6060001683e0, value=57 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=58, fence=0x604000895650 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600086a9e0, from_fence=0x6060008c8700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c8dc0, semaphore=0x6060001683e0, value=58 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0007d79a0, f=0, wait_fence=0x60600086a9e0 {0x6060001683e0:57}, signal_fence=0x604000895650 {0x6060001683e0:58} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006d96a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600035b820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006d96a0, semaphore=0x6060001683e0, value=58 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600035b820, semaphore=0x6060001683e0, value=58 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600086a680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600086a680, semaphore=0x6060001683e0, value=58 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=59, fence=0x60400163e550 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600086a680, from_fence=0x6060008c5820 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d6e5c0, semaphore=0x6060001683e0, value=59 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0003d73a0, f=0, wait_fence=0x60600086a680 {0x6060001683e0:58}, signal_fence=0x60400163e550 {0x6060001683e0:59} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600086a5c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000cbc960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600086a5c0, semaphore=0x6060001683e0, value=59 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cbc960, semaphore=0x6060001683e0, value=59 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600011b4c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600011b4c0, semaphore=0x6060001683e0, value=59 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=60, fence=0x60400154e050 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600011b4c0, from_fence=0x606000d4b580 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000254d20, semaphore=0x6060001683e0, value=60 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x60600011b4c0 {0x6060001683e0:59}, signal_fence=0x60400154e050 {0x6060001683e0:60} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006fdf40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000452720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006fdf40, semaphore=0x6060001683e0, value=60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000452720, semaphore=0x6060001683e0, value=60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000791540 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000791540, semaphore=0x6060001683e0, value=60 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=61, fence=0x6040000eaa50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000791540, from_fence=0x60600036e060 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072eae0, semaphore=0x6060001683e0, value=61 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000843b60, f=0, wait_fence=0x606000791540 {0x6060001683e0:60}, signal_fence=0x6040000eaa50 {0x6060001683e0:61} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c07160 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003f8480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c07160, semaphore=0x6060001683e0, value=61 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003f8480, semaphore=0x6060001683e0, value=61 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003f83c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003f83c0, semaphore=0x6060001683e0, value=61 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=62, fence=0x604000e02990 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003f83c0, from_fence=0x606000129680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001cd8a0, semaphore=0x6060001683e0, value=62 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x6060003f83c0 {0x6060001683e0:61}, signal_fence=0x604000e02990 {0x6060001683e0:62} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009a7960 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d16de0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009a7960, semaphore=0x6060001683e0, value=62 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d16de0, semaphore=0x6060001683e0, value=62 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003e7500 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003e7500, semaphore=0x6060001683e0, value=62 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=63, fence=0x604001627210 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060003e7500, from_fence=0x6060000a1780 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600052e1a0, semaphore=0x6060001683e0, value=63 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x6060003e7500 {0x6060001683e0:62}, signal_fence=0x604001627210 {0x6060001683e0:63} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005cd560 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007c6280 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005cd560, semaphore=0x6060001683e0, value=63 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007c6280, semaphore=0x6060001683e0, value=63 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004c9e20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004c9e20, semaphore=0x6060001683e0, value=63 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=64, fence=0x604000d12d50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060004c9e20, from_fence=0x606000ccbcc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072c140, semaphore=0x6060001683e0, value=64 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x6060004c9e20 {0x6060001683e0:63}, signal_fence=0x604000d12d50 {0x6060001683e0:64} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002c5ee0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c5340 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002c5ee0, semaphore=0x6060001683e0, value=64 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c5340, semaphore=0x6060001683e0, value=64 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b40ec0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b40ec0, semaphore=0x6060001683e0, value=64 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=65, fence=0x604001510cd0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000b40ec0, from_fence=0x60600036fbc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008daca0, semaphore=0x6060001683e0, value=65 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00070d380, f=0, wait_fence=0x606000b40ec0 {0x6060001683e0:64}, signal_fence=0x604001510cd0 {0x6060001683e0:65} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600038bd00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b5d360 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600038bd00, semaphore=0x6060001683e0, value=65 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b5d360, semaphore=0x6060001683e0, value=65 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000879140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000879140, semaphore=0x6060001683e0, value=65 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=66, fence=0x604000f3c8d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000879140, from_fence=0x606000b37da0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b5400, semaphore=0x6060001683e0, value=66 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a36100, f=0, wait_fence=0x606000879140 {0x6060001683e0:65}, signal_fence=0x604000f3c8d0 {0x6060001683e0:66} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006d6160 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a045a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006d6160, semaphore=0x6060001683e0, value=66 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a045a0, semaphore=0x6060001683e0, value=66 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600062edc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600062edc0, semaphore=0x6060001683e0, value=66 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=67, fence=0x604000533b50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600062edc0, from_fence=0x606000283ac0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003d5aa0, semaphore=0x6060001683e0, value=67 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00036c520, f=0, wait_fence=0x60600062edc0 {0x6060001683e0:66}, signal_fence=0x604000533b50 {0x6060001683e0:67} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005686a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b56e20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005686a0, semaphore=0x6060001683e0, value=67 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b56e20, semaphore=0x6060001683e0, value=67 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005b4d20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005b4d20, semaphore=0x6060001683e0, value=67 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=68, fence=0x60400166c610 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005b4d20, from_fence=0x6060005b2e60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600073d120, semaphore=0x6060001683e0, value=68 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0003c8250, f=0, wait_fence=0x6060005b4d20 {0x6060001683e0:67}, signal_fence=0x60400166c610 {0x6060001683e0:68} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a1fae0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b57240 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a1fae0, semaphore=0x6060001683e0, value=68 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b57240, semaphore=0x6060001683e0, value=68 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a3a780 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a3a780, semaphore=0x6060001683e0, value=68 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=69, fence=0x604000f0e110 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000a3a780, from_fence=0x606000bccc60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c7aa20, semaphore=0x6060001683e0, value=69 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00011f9b0, f=0, wait_fence=0x606000a3a780 {0x6060001683e0:68, 0x6060001682c0:92}, signal_fence=0x604000f0e110 {0x6060001683e0:69} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000abce00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000771f80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000abce00, semaphore=0x6060001683e0, value=69 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000771f80, semaphore=0x6060001683e0, value=69 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0006553c0, wait={0x6060001682c0:92, 0x6060001683e0:69}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b731c0, wait={0x6060001683e0:68}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00023b6c0, wait={0x6060001683e0:67}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009b26c0, wait={0x6060001683e0:66}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000354dc0, wait={0x6060001683e0:65}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000a85ec0, wait={0x6060001683e0:64}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000320b00, wait={0x6060001683e0:63}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00014f200, wait={0x6060001683e0:62}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000406000, wait={0x6060001683e0:61}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00054f4c0, wait={0x6060001683e0:60}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00052e700, wait={0x6060001683e0:59}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b473c0, wait={0x6060001683e0:58}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00035e000, wait={0x6060001683e0:57}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b13280, wait={0x6060001683e0:56}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009b2840, wait={0x6060001683e0:55}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000bc53c0, wait={0x6060001683e0:54}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00041d700, wait={0x6060001683e0:53}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000319540, wait={0x6060001683e0:52}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b2f480, wait={0x6060001683e0:51}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0005eb640, wait={0x6060001683e0:50}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000bc7e80, wait={0x6060001683e0:49}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00052ed00, wait={0x6060001683e0:48}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000639880, wait={0x6060001683e0:47}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000b72c80, wait={0x6060001682c0:90, 0x6060001683e0:46}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0000a2c40, wait={0x6060001682c0:44, 0x6060001683e0:45}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00034f540, wait={0x6060001682c0:42, 0x6060001683e0:44}, signal={} (OK)
W0605 17:37:47.177995 139756486557120 dispatch.py:272] Finished tracing + transforming jit(convert_element_type) in 0.00024437904357910156 sec
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000243a00 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:92}, signal={0x6060001682c0:93} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:93}, signal={0x6060001682c0:94} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a17c20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004cfdc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a17c20, semaphore=0x6060001682c0, value=94 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004cfdc0, semaphore=0x6060001682c0, value=94 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x6080025003a0, wait={0x6060001683e0:10}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x608002500320, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x6080025002a0, wait={0x6060001683e0:10}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x608002500220, wait={0x6060001683e0:10}, signal={} (OK)
I0605 17:37:47.196380 139756486557120 partitioning.py:631] train state shapes: TrainState(step=(), mdl_vars={'params': {'lm': {'final_ln': {'bias': (2048,), 'scale': (2048,)}, 'position_emb': {'emb_var': (2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (51200,)}, 'linear': {'w': (2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (24, 8192)}, 'linear': {'w': (24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (24, 2048)}, 'linear': {'w': (24, 8192, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}, 'self_attention': {'combined_qkv': {'w': (24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (24, 64)}, 'post': {'w': (24, 2048, 32, 64)}}}}}}}}}, opt_states=[{'no_prefix': ({'count': ()}, {'count': ()}, {'count': (), 'm': {'params': {'lm': {'final_ln': {'bias': (2048,), 'scale': (2048,)}, 'position_emb': {'emb_var': (2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (51200,)}, 'linear': {'w': (2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'ffn_layer2': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'self_attention': {'combined_qkv': {'w': MaskedNode()}, 'per_dim_scale': {'per_dim_scale': MaskedNode()}, 'post': {'w': MaskedNode()}}}}}}}}}, 'v': {'params': {'lm': {'final_ln': {'bias': (2048,), 'scale': (2048,)}, 'position_emb': {'emb_var': (2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (51200,)}, 'linear': {'w': (2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'ffn_layer2': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'self_attention': {'combined_qkv': {'w': MaskedNode()}, 'per_dim_scale': {'per_dim_scale': MaskedNode()}, 'post': {'w': MaskedNode()}}}}}}}}}}, {'count': ()}), 'p#24#i-1': ({'count': (24,)}, {'count': (24,)}, {'count': (24,), 'm': {'params': {'lm': {'final_ln': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'position_emb': {'emb_var': MaskedNode()}, 'softmax': {'logits_ffn': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (24, 8192)}, 'linear': {'w': (24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (24, 2048)}, 'linear': {'w': (24, 8192, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}, 'self_attention': {'combined_qkv': {'w': (24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (24, 64)}, 'post': {'w': (24, 2048, 32, 64)}}}}}}}}}, 'v': {'params': {'lm': {'final_ln': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'position_emb': {'emb_var': MaskedNode()}, 'softmax': {'logits_ffn': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (24, 8192)}, 'linear': {'w': (24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (24, 2048)}, 'linear': {'w': (24, 8192, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}}, 'layer_norm': {'bias': (24, 2048), 'scale': (24, 2048)}, 'self_attention': {'combined_qkv': {'w': (24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (24, 64)}, 'post': {'w': (24, 2048, 32, 64)}}}}}}}}}}, {'count': (24,)})}])
I0605 17:37:47.209682 139756486557120 partitioning.py:637] replicated train state shapes: TrainState(step=(1,), mdl_vars={'params': {'lm': {'final_ln': {'bias': (1, 2048), 'scale': (1, 2048)}, 'position_emb': {'emb_var': (1, 2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (1, 51200)}, 'linear': {'w': (1, 2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (1, 24, 8192)}, 'linear': {'w': (1, 24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (1, 24, 2048)}, 'linear': {'w': (1, 24, 8192, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}, 'self_attention': {'combined_qkv': {'w': (1, 24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (1, 24, 64)}, 'post': {'w': (1, 24, 2048, 32, 64)}}}}}}}}}, opt_states=[{'no_prefix': ({'count': (1,)}, {'count': (1,)}, {'count': (1,), 'm': {'params': {'lm': {'final_ln': {'bias': (1, 2048), 'scale': (1, 2048)}, 'position_emb': {'emb_var': (1, 2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (1, 51200)}, 'linear': {'w': (1, 2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'ffn_layer2': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'self_attention': {'combined_qkv': {'w': MaskedNode()}, 'per_dim_scale': {'per_dim_scale': MaskedNode()}, 'post': {'w': MaskedNode()}}}}}}}}}, 'v': {'params': {'lm': {'final_ln': {'bias': (1, 2048), 'scale': (1, 2048)}, 'position_emb': {'emb_var': (1, 2048, 2048)}, 'softmax': {'logits_ffn': {'bias': {'b': (1, 51200)}, 'linear': {'w': (1, 2048, 51200)}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'ffn_layer2': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}}, 'layer_norm': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'self_attention': {'combined_qkv': {'w': MaskedNode()}, 'per_dim_scale': {'per_dim_scale': MaskedNode()}, 'post': {'w': MaskedNode()}}}}}}}}}}, {'count': (1,)}), 'p#24#i-1': ({'count': (1, 24)}, {'count': (1, 24)}, {'count': (1, 24), 'm': {'params': {'lm': {'final_ln': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'position_emb': {'emb_var': MaskedNode()}, 'softmax': {'logits_ffn': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (1, 24, 8192)}, 'linear': {'w': (1, 24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (1, 24, 2048)}, 'linear': {'w': (1, 24, 8192, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}, 'self_attention': {'combined_qkv': {'w': (1, 24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (1, 24, 64)}, 'post': {'w': (1, 24, 2048, 32, 64)}}}}}}}}}, 'v': {'params': {'lm': {'final_ln': {'bias': MaskedNode(), 'scale': MaskedNode()}, 'position_emb': {'emb_var': MaskedNode()}, 'softmax': {'logits_ffn': {'bias': {'b': MaskedNode()}, 'linear': {'w': MaskedNode()}}}, 'transformer': {'repeat': {'sub': {'x_layers_0': {'ff_layer': {'ffn_layer1': {'bias': {'b': (1, 24, 8192)}, 'linear': {'w': (1, 24, 2048, 8192)}}, 'ffn_layer2': {'bias': {'b': (1, 24, 2048)}, 'linear': {'w': (1, 24, 8192, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}}, 'layer_norm': {'bias': (1, 24, 2048), 'scale': (1, 24, 2048)}, 'self_attention': {'combined_qkv': {'w': (1, 24, 3, 2048, 32, 64)}, 'per_dim_scale': {'per_dim_scale': (1, 24, 64)}, 'post': {'w': (1, 24, 2048, 32, 64)}}}}}}}}}}, {'count': (1, 24)})}])
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000c53080 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:94}, signal={0x6060001682c0:95} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:95}, signal={0x6060001682c0:96} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000882020 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c3f200 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000882020, semaphore=0x6060001682c0, value=96 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c3f200, semaphore=0x6060001682c0, value=96 (OK)
W0605 17:37:47.211523 139756486557120 pxla.py:1882] Compiling _threefry_fold_in for with global shapes and types [ShapedArray(uint32[2]), ShapedArray(uint32[])]. Argument mapping: (GSPMDSharding({replicated}), GSPMDSharding({replicated})).
W0605 17:37:47.269204 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_fold_in) in 0.05749940872192383 sec
W0605 17:37:48.786752 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_fold_in) in 1.5170984268188477 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600042b8a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600042b8a0, semaphore=0x6060001683e0, value=69 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=70, fence=0x60400042c5d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600042b8a0, from_fence=0x606000328ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003292a0, semaphore=0x6060001683e0, value=70 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600042b8a0, from_fence=0x606000882020 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c3f200, semaphore=0x6060001683e0, value=70 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0008228a0, f=0, wait_fence=0x60600042b8a0 {0x6060001683e0:69, 0x6060001682c0:96}, signal_fence=0x60400042c5d0 {0x6060001683e0:70} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006155c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600014e580 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006155c0, semaphore=0x6060001683e0, value=70 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600014e580, semaphore=0x6060001683e0, value=70 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000c54580, wait={0x6060001682c0:96, 0x6060001683e0:70}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x608001c844a0, wait={0x6060001683e0:70}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc14180 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:70}, signal={} (OK)
I0605 17:37:48.789292 139756486557120 partitioning.py:647] root prng key: [3199903509 2250625448]
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x608001c84420, wait={0x6060001683e0:9}, signal={} (OK)
W0605 17:37:48.797018 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00022339820861816406 sec
W0605 17:37:48.797891 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0017151832580566406 sec
W0605 17:37:48.798967 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.003200054168701172 sec
W0605 17:37:48.799690 139756486557120 pxla.py:1882] Compiling _threefry_split_original for with global shapes and types [ShapedArray(uint32[2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:48.992391 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003879070281982422 sec
W0605 17:37:48.993552 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003216266632080078 sec
W0605 17:37:48.994482 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003116130828857422 sec
W0605 17:37:48.995203 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003108978271484375 sec
W0605 17:37:49.045017 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_split_original) in 0.24518394470214844 sec
W0605 17:37:51.420535 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_split_original) in 2.3751134872436523 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a76ba0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a76ba0, semaphore=0x6060001683e0, value=70 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=71, fence=0x604000dd4f50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000a76ba0, from_fence=0x6060006155c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600014e580, semaphore=0x6060001683e0, value=71 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000a10140, f=0, wait_fence=0x606000a76ba0 {0x6060001683e0:70}, signal_fence=0x604000dd4f50 {0x6060001683e0:71} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006b5160 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b37980 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b5160, semaphore=0x6060001683e0, value=71 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b37980, semaphore=0x6060001683e0, value=71 (OK)
W0605 17:37:51.424489 139756486557120 dispatch.py:272] Finished tracing + transforming _unstack for pjit in 0.001565694808959961 sec
W0605 17:37:51.425680 139756486557120 pxla.py:1882] Compiling _unstack for with global shapes and types [ShapedArray(uint32[3,2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:51.432903 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_unstack) in 0.006999969482421875 sec
W0605 17:37:51.687622 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_unstack) in 0.25420069694519043 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060000a7300 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060000a7300, semaphore=0x6060001683e0, value=71 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=72, fence=0x60400156c090 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060000a7300, from_fence=0x6060006b5160 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b37980, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000842820, f=0, wait_fence=0x6060000a7300 {0x6060001683e0:71}, signal_fence=0x60400156c090 {0x6060001683e0:72} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000af6700 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003b2ca0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000af6700, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003b2ca0, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c003e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009390e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c003e0, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009390e0, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600037fd00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600038e700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600037fd00, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600038e700, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000387e80, wait={0x6060001683e0:72}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc14520 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:72}, signal={} (OK)
I0605 17:37:51.688758 139756486557120 executors.py:260] train prng seed: [3373580220 3771856083]
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc14520 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:72}, signal={} (OK)
I0605 17:37:51.689568 139756486557120 executors.py:261] eval prng seed: [3893388808  331134876]
W0605 17:37:51.691859 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.0013885498046875 sec
W0605 17:37:51.692612 139756486557120 pxla.py:1882] Compiling _threefry_split_original for with global shapes and types [ShapedArray(uint32[2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:37:51.749565 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_threefry_split_original) in 0.05680084228515625 sec
W0605 17:37:53.214284 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_threefry_split_original) in 1.4642627239227295 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c8c3c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c8c3c0, semaphore=0x6060001683e0, value=72 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=73, fence=0x604001124f50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c8c3c0, from_fence=0x606000c003e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009390e0, semaphore=0x6060001683e0, value=73 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0005ea050, f=0, wait_fence=0x606000c8c3c0 {0x6060001683e0:72}, signal_fence=0x604001124f50 {0x6060001683e0:73} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000359e40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008a3080 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000359e40, semaphore=0x6060001683e0, value=73 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a3080, semaphore=0x6060001683e0, value=73 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x608002259820, wait={0x6060001683e0:73}, signal={} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008c2160 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c2160, semaphore=0x6060001683e0, value=73 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=74, fence=0x604001125cd0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060008c2160, from_fence=0x60600037fd00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600038e700, semaphore=0x6060001683e0, value=74 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0005ea050, f=0, wait_fence=0x6060008c2160 {0x6060001683e0:73}, signal_fence=0x604001125cd0 {0x6060001683e0:74} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008a31a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000abf920 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a31a0, semaphore=0x6060001683e0, value=74 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000abf920, semaphore=0x6060001683e0, value=74 (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60800397ffa0, wait={0x6060001683e0:74}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00029f140, wait={0x6060001683e0:71}, signal={} (OK)
I0605 17:37:53.216853 139756486557120 executors.py:295] Starting executor.
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc160c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:94}, signal={} (OK)
I0605 17:37:53.217641 139756486557120 executors.py:454] Model initial global_step=0
I0605 17:37:53.217703 139756486557120 executors.py:461] [PAX STATUS]: Starting training loop.
I0605 17:37:53.217766 139756486557120 programs.py:210] [PAX STATUS]: Setting up BaseTrainProgram.
I0605 17:37:53.217862 139756486557120 summary_utils.py:281] Opening SummaryWriter `log_NVIDIA1_3BPmap/summaries/train`...
I0605 17:37:53.219248 139756486557120 summary_utils.py:281] Opening SummaryWriter `log_NVIDIA1_3BPmap/summaries/eval_train`...
I0605 17:37:53.226810 139756486557120 py_utils.py:338] Starting sync_global_devices Start training loop from step: 0 across 1 devices globally
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00084f880 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:96}, signal={0x6060001682c0:97} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:97}, signal={0x6060001682c0:98} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600052c760 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000189920 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600052c760, semaphore=0x6060001682c0, value=98 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000189920, semaphore=0x6060001682c0, value=98 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a11320 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a11320, semaphore=0x6060001683e0, value=74 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=75, fence=0x60400056c590 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000a11320, from_fence=0x60600052c760 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000189920, semaphore=0x6060001683e0, value=75 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x606000a11320 {0x6060001683e0:74, 0x6060001682c0:98}, signal_fence=0x60400056c590 {0x6060001683e0:75} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004e58a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008859e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004e58a0, semaphore=0x6060001683e0, value=75 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008859e0, semaphore=0x6060001683e0, value=75 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc149e0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:75}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00084f7c0, wait={0x6060001682c0:98, 0x6060001683e0:75}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00084f7c0, wait={0x6060001683e0:75}, signal={} (OK)
I0605 17:37:53.229556 139756486557120 py_utils.py:341] Finished sync_global_devices Start training loop from step: 0 across 1 devices globally
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b2a080 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:98}, signal={0x6060001682c0:99} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:99}, signal={0x6060001682c0:100} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003aa000 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a9ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003aa000, semaphore=0x6060001682c0, value=100 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a9ee0, semaphore=0x6060001682c0, value=100 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0006aa380 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:100}, signal={0x6060001682c0:101} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:101}, signal={0x6060001682c0:102} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600047ea60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008a2780 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600047ea60, semaphore=0x6060001682c0, value=102 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a2780, semaphore=0x6060001682c0, value=102 (OK)
W0605 17:37:53.425011 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0005414485931396484 sec
W0605 17:37:53.425846 139756486557120 dispatch.py:272] Finished tracing + transforming _psum for pjit in 0.0017685890197753906 sec
W0605 17:37:53.426665 139756486557120 pxla.py:1882] Compiling _psum for with global shapes and types [ShapedArray(int32[1]), ShapedArray(int32[1])]. Argument mapping: (GSPMDSharding({maximal device=0}), GSPMDSharding({maximal device=0})).
W0605 17:37:53.432121 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_psum) in 0.0052835941314697266 sec
W0605 17:37:53.736919 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_psum) in 0.3044319152832031 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002cf660 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002cf660, semaphore=0x6060001683e0, value=75 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=76, fence=0x6040016b70d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060002cf660, from_fence=0x6060003aa000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a9ee0, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060002cf660, from_fence=0x60600047ea60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a2780, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004abaa0, f=0, wait_fence=0x6060002cf660 {0x6060001683e0:75, 0x6060001682c0:102}, signal_fence=0x6040016b70d0 {0x6060001683e0:76} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ad7740 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c95720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ad7740, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c95720, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600031bf20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004fb560 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600031bf20, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004fb560, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13bc0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13c00 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000912dc0, wait={0x6060001682c0:102, 0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000ca0d80, wait={0x6060001682c0:100, 0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000912dc0, wait={0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000ca0d80, wait={0x6060001683e0:76}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000285640 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:102}, signal={0x6060001682c0:103} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:103}, signal={0x6060001682c0:104} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a3b380 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005b2260 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a3b380, semaphore=0x6060001682c0, value=104 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005b2260, semaphore=0x6060001682c0, value=104 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004084a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004084a0, semaphore=0x6060001683e0, value=76 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=77, fence=0x604000e0f4d0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060004084a0, from_fence=0x606000a3b380 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005b2260, semaphore=0x6060001683e0, value=77 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060004084a0 {0x6060001683e0:76, 0x6060001682c0:104}, signal_fence=0x604000e0f4d0 {0x6060001683e0:77} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000845b40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b04c20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000845b40, semaphore=0x6060001683e0, value=77 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b04c20, semaphore=0x6060001683e0, value=77 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc139e0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:77}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000285400, wait={0x6060001682c0:104, 0x6060001683e0:77}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000285400, wait={0x6060001683e0:77}, signal={} (OK)
I0605 17:37:53.741610 139756486557120 checkpointer.py:67] Saving item to log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/state.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00032d880 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:104}, signal={0x6060001682c0:105} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:105}, signal={0x6060001682c0:106} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006963e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a81220 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006963e0, semaphore=0x6060001682c0, value=106 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a81220, semaphore=0x6060001682c0, value=106 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000aab3c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:106}, signal={0x6060001682c0:107} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:107}, signal={0x6060001682c0:108} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005cf000 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008a9e00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005cf000, semaphore=0x6060001682c0, value=108 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a9e00, semaphore=0x6060001682c0, value=108 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b9d800 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b9d800, semaphore=0x6060001683e0, value=77 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=78, fence=0x6040013f2a90 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000b9d800, from_fence=0x6060006963e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a81220, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000b9d800, from_fence=0x6060005cf000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a9e00, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004abaa0, f=0, wait_fence=0x606000b9d800 {0x6060001683e0:77, 0x6060001682c0:108}, signal_fence=0x6040013f2a90 {0x6060001683e0:78} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000603080 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000603560 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000603080, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000603560, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060009b2d60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005f2c40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009b2d60, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005f2c40, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d40 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d80 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000116500, wait={0x6060001682c0:108, 0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0003b5240, wait={0x6060001682c0:106, 0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000116500, wait={0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0003b5240, wait={0x6060001683e0:78}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000843040 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:108}, signal={0x6060001682c0:109} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:109}, signal={0x6060001682c0:110} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a94000 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007b61a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a94000, semaphore=0x6060001682c0, value=110 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007b61a0, semaphore=0x6060001682c0, value=110 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600092a740 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600092a740, semaphore=0x6060001683e0, value=78 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=79, fence=0x6040014e0250 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600092a740, from_fence=0x606000a94000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007b61a0, semaphore=0x6060001683e0, value=79 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x60600092a740 {0x6060001683e0:78, 0x6060001682c0:110}, signal_fence=0x6040014e0250 {0x6060001683e0:79} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003d4060 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000910dc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003d4060, semaphore=0x6060001683e0, value=79 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000910dc0, semaphore=0x6060001683e0, value=79 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13b60 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:79}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000842ec0, wait={0x6060001682c0:110, 0x6060001683e0:79}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000842ec0, wait={0x6060001683e0:79}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:94}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:11}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:14}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:16}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:38}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:12}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:14}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:16}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:18}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:20}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:13}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:15}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:17}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:19}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:21}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:40}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:44}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:45}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:46}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:47}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:48}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:49}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:50}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:51}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:52}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:53}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:54}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:55}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:56}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:57}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:58}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:59}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:60}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:61}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:62}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:63}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:64}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:65}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:66}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:67}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:68}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc148c0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:69}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0005a61c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:110}, signal={0x6060001682c0:111} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:111}, signal={0x6060001682c0:112} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004a1260 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600049cee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004a1260, semaphore=0x6060001682c0, value=112 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600049cee0, semaphore=0x6060001682c0, value=112 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600008f2a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600008f2a0, semaphore=0x6060001683e0, value=79 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=80, fence=0x6040014e0fd0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600008f2a0, from_fence=0x6060004a1260 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600049cee0, semaphore=0x6060001683e0, value=80 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x60600008f2a0 {0x6060001683e0:79, 0x6060001682c0:112}, signal_fence=0x6040014e0fd0 {0x6060001683e0:80} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000715580 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008ae960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000715580, semaphore=0x6060001683e0, value=80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008ae960, semaphore=0x6060001683e0, value=80 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc12a60 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:80}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000840040, wait={0x6060001682c0:112, 0x6060001683e0:80}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000840040, wait={0x6060001683e0:80}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0008c2180 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:112}, signal={0x6060001682c0:113} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:113}, signal={0x6060001682c0:114} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060001ad4a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004fb020 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001ad4a0, semaphore=0x6060001682c0, value=114 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004fb020, semaphore=0x6060001682c0, value=114 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c65fc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c65fc0, semaphore=0x6060001683e0, value=80 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=81, fence=0x6040014e0fd0 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000c65fc0, from_fence=0x6060001ad4a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004fb020, semaphore=0x6060001683e0, value=81 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x606000c65fc0 {0x6060001683e0:80, 0x6060001682c0:114}, signal_fence=0x6040014e0fd0 {0x6060001683e0:81} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b320a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b32760 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b320a0, semaphore=0x6060001683e0, value=81 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b32760, semaphore=0x6060001683e0, value=81 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13ac0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:81}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0008c2240, wait={0x6060001682c0:114, 0x6060001683e0:81}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0008c2240, wait={0x6060001683e0:81}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000887140 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:114}, signal={0x6060001682c0:115} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:115}, signal={0x6060001682c0:116} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000971c00 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000971f00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000971c00, semaphore=0x6060001682c0, value=116 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000971f00, semaphore=0x6060001682c0, value=116 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000673be0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000673be0, semaphore=0x6060001683e0, value=81 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=82, fence=0x6040003e5f10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x606000673be0, from_fence=0x606000971c00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000971f00, semaphore=0x6060001683e0, value=82 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x606000673be0 {0x6060001683e0:81, 0x6060001682c0:116}, signal_fence=0x6040003e5f10 {0x6060001683e0:82} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004c0b20 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000604880 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004c0b20, semaphore=0x6060001683e0, value=82 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000604880, semaphore=0x6060001683e0, value=82 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d40 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:82}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000887080, wait={0x6060001682c0:116, 0x6060001683e0:82}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000887080, wait={0x6060001683e0:82}, signal={} (OK)
I0605 17:38:48.673480 139756486557120 utils.py:465] Renaming log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/state.orbax-checkpoint-tmp-1685986673741760 to log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/state
I0605 17:38:48.673784 139756486557120 utils.py:509] Finished saving checkpoint to `log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/state`.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b785c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:116}, signal={0x6060001682c0:117} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:117}, signal={0x6060001682c0:118} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600096d6a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000875f00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600096d6a0, semaphore=0x6060001682c0, value=118 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000875f00, semaphore=0x6060001682c0, value=118 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006c6980 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006c6980, semaphore=0x6060001683e0, value=82 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=83, fence=0x60400042e250 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006c6980, from_fence=0x60600096d6a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000875f00, semaphore=0x6060001683e0, value=83 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060006c6980 {0x6060001683e0:82, 0x6060001682c0:118}, signal_fence=0x60400042e250 {0x6060001683e0:83} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600062b280 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007c6460 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600062b280, semaphore=0x6060001683e0, value=83 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007c6460, semaphore=0x6060001683e0, value=83 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d40 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:83}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000390580, wait={0x6060001682c0:118, 0x6060001683e0:83}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000390580, wait={0x6060001683e0:83}, signal={} (OK)
I0605 17:38:48.676459 139756486557120 checkpointer.py:67] Saving item to log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/metadata.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00054f100 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:118}, signal={0x6060001682c0:119} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:119}, signal={0x6060001682c0:120} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004b38c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000375680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004b38c0, semaphore=0x6060001682c0, value=120 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000375680, semaphore=0x6060001682c0, value=120 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c0008cff80 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:120}, signal={0x6060001682c0:121} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:121}, signal={0x6060001682c0:122} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000256ca0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600065abc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000256ca0, semaphore=0x6060001682c0, value=122 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600065abc0, semaphore=0x6060001682c0, value=122 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600059b460 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600059b460, semaphore=0x6060001683e0, value=83 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=84, fence=0x604000858590 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600059b460, from_fence=0x6060004b38c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000375680, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600059b460, from_fence=0x606000256ca0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600065abc0, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b0004abaa0, f=0, wait_fence=0x60600059b460 {0x6060001683e0:83, 0x6060001682c0:122}, signal_fence=0x604000858590 {0x6060001683e0:84} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000d3b0e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060002890a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d3b0e0, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002890a0, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a588a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000788960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a588a0, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000788960, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d80 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13dc0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0008cfec0, wait={0x6060001682c0:122, 0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0007d9140, wait={0x6060001682c0:120, 0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0008cfec0, wait={0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0007d9140, wait={0x6060001683e0:84}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00025b640 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:122}, signal={0x6060001682c0:123} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:123}, signal={0x6060001682c0:124} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005eb320 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060004427c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005eb320, semaphore=0x6060001682c0, value=124 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004427c0, semaphore=0x6060001682c0, value=124 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060005237c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005237c0, semaphore=0x6060001683e0, value=84 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=85, fence=0x6040005b3310 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060005237c0, from_fence=0x6060005eb320 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004427c0, semaphore=0x6060001683e0, value=85 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060005237c0 {0x6060001683e0:84, 0x6060001682c0:124}, signal_fence=0x6040005b3310 {0x6060001683e0:85} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060003a0820 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b66720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a0820, semaphore=0x6060001683e0, value=85 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b66720, semaphore=0x6060001683e0, value=85 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13bc0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:85}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00025bb80, wait={0x6060001682c0:124, 0x6060001683e0:85}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00025bb80, wait={0x6060001683e0:85}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c00098fe00 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:124}, signal={0x6060001682c0:125} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:125}, signal={0x6060001682c0:126} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600070c5e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006cc260 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600070c5e0, semaphore=0x6060001682c0, value=126 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006cc260, semaphore=0x6060001682c0, value=126 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600078b900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600078b900, semaphore=0x6060001683e0, value=85 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=86, fence=0x604001016790 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600078b900, from_fence=0x60600070c5e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006cc260, semaphore=0x6060001683e0, value=86 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x60600078b900 {0x6060001683e0:85, 0x6060001682c0:126}, signal_fence=0x604001016790 {0x6060001683e0:86} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000736e80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000536780 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000736e80, semaphore=0x6060001683e0, value=86 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000536780, semaphore=0x6060001683e0, value=86 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13b60 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:86}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00098fec0, wait={0x6060001682c0:126, 0x6060001683e0:86}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c00098fec0, wait={0x6060001683e0:86}, signal={} (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000769780 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:126}, signal={0x6060001682c0:127} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:127}, signal={0x6060001682c0:128} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007da560 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000cf62c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007da560, semaphore=0x6060001682c0, value=128 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cf62c0, semaphore=0x6060001682c0, value=128 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600059a1a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600059a1a0, semaphore=0x6060001683e0, value=86 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=87, fence=0x604001515d10 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600059a1a0, from_fence=0x6060007da560 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cf62c0, semaphore=0x6060001683e0, value=87 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x60600059a1a0 {0x6060001683e0:86, 0x6060001682c0:128}, signal_fence=0x604001515d10 {0x6060001683e0:87} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007b2e40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060007c4420 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007b2e40, semaphore=0x6060001683e0, value=87 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007c4420, semaphore=0x6060001683e0, value=87 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d80 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:87}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000769600, wait={0x6060001682c0:128, 0x6060001683e0:87}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c000769600, wait={0x6060001683e0:87}, signal={} (OK)
I0605 17:38:48.687450 139756486557120 utils.py:465] Renaming log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/metadata.orbax-checkpoint-tmp-1685986728676589 to log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/metadata
I0605 17:38:48.687598 139756486557120 utils.py:509] Finished saving checkpoint to `log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257/metadata`.
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000b5b340 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:128}, signal={0x6060001682c0:129} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:129}, signal={0x6060001682c0:130} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000330e60 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a94000 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000330e60, semaphore=0x6060001682c0, value=130 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a94000, semaphore=0x6060001682c0, value=130 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060008dd7c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008dd7c0, semaphore=0x6060001683e0, value=87 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=88, fence=0x604000f7d010 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060008dd7c0, from_fence=0x606000330e60 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a94000, semaphore=0x6060001683e0, value=88 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060008dd7c0 {0x6060001683e0:87, 0x6060001682c0:130}, signal_fence=0x604000f7d010 {0x6060001683e0:88} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600087c7a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000ccc140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600087c7a0, semaphore=0x6060001683e0, value=88 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000ccc140, semaphore=0x6060001683e0, value=88 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13d80 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:88}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0002a1b40, wait={0x6060001682c0:130, 0x6060001683e0:88}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0002a1b40, wait={0x6060001683e0:88}, signal={} (OK)
I0605 17:38:48.689893 139756486557120 utils.py:465] Renaming log_NVIDIA1_3BPmap/checkpoints/checkpoint_0.orbax-checkpoint-tmp-1685986673422257 to log_NVIDIA1_3BPmap/checkpoints/checkpoint_0
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c000282d00 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:130}, signal={0x6060001682c0:131} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:131}, signal={0x6060001682c0:132} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000a0a4e0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000486c20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a0a4e0, semaphore=0x6060001682c0, value=132 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000486c20, semaphore=0x6060001682c0, value=132 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060006c8ba0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006c8ba0, semaphore=0x6060001683e0, value=88 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=89, fence=0x604000600650 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x6060006c8ba0, from_fence=0x606000a0a4e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000486c20, semaphore=0x6060001683e0, value=89 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b00038c290, f=0, wait_fence=0x6060006c8ba0 {0x6060001683e0:88, 0x6060001682c0:132}, signal_fence=0x604000600650 {0x6060001683e0:89} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b9d8c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000b9db00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b9d8c0, semaphore=0x6060001683e0, value=89 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b9db00, semaphore=0x6060001683e0, value=89 (OK)
:: IREE INVOKE (hal_allocator_import_buffer): external_buffer=0x7ffdadc13fc0 (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001683e0:89}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009bfbc0, wait={0x6060001682c0:132, 0x6060001683e0:89}, signal={} (OK)
:: IREE INVOKE (hal_device_queue_dealloca): device=0x6110011f9800, buffer=0x60c0009bfbc0, wait={0x6060001683e0:89}, signal={} (OK)
W0605 17:38:50.088575 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00046539306640625 sec
W0605 17:38:50.089799 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004601478576660156 sec
W0605 17:38:50.090771 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003387928009033203 sec
I0605 17:38:50.104525 139756486557120 base_layer.py:632] Creating var /lm/softmax/logits_ffn/linear/w with shape=[2048, 51200], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.022097086912079608
W0605 17:38:50.106836 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004742145538330078 sec
W0605 17:38:50.107717 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003573894500732422 sec
W0605 17:38:50.108556 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003871917724609375 sec
W0605 17:38:50.109402 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038361549377441406 sec
W0605 17:38:50.110023 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.004520893096923828 sec
W0605 17:38:50.110689 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.005500078201293945 sec
W0605 17:38:50.111018 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.006124973297119141 sec
W0605 17:38:50.111812 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004010200500488281 sec
W0605 17:38:50.115243 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00039196014404296875 sec
W0605 17:38:50.116189 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00017595291137695312 sec
W0605 17:38:50.117643 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00039315223693847656 sec
I0605 17:38:50.120159 139756486557120 base_layer.py:632] Creating var /lm/position_emb/emb_var with shape=[2048, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
W0605 17:38:50.122282 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.000461578369140625 sec
W0605 17:38:50.123463 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003762245178222656 sec
W0605 17:38:50.124331 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038623809814453125 sec
W0605 17:38:50.124928 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.0038928985595703125 sec
W0605 17:38:50.125589 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.004845380783081055 sec
W0605 17:38:50.125927 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.005454301834106445 sec
W0605 17:38:50.126752 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00040984153747558594 sec
W0605 17:38:50.131844 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00037741661071777344 sec
W0605 17:38:50.132391 139756486557120 dispatch.py:272] Finished tracing + transforming _one_hot for pjit in 0.0014386177062988281 sec
W0605 17:38:50.133401 139756486557120 dispatch.py:272] Finished tracing + transforming matmul for pjit in 0.0005931854248046875 sec
W0605 17:38:50.136929 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006892681121826172 sec
W0605 17:38:50.139818 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003864765167236328 sec
W0605 17:38:50.142343 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003662109375 sec
W0605 17:38:50.143554 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036787986755371094 sec
W0605 17:38:50.146310 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003838539123535156 sec
W0605 17:38:50.147240 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003628730773925781 sec
W0605 17:38:50.148331 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003795623779296875 sec
W0605 17:38:50.226294 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004551410675048828 sec
W0605 17:38:50.227395 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00037407875061035156 sec
W0605 17:38:50.228480 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004832744598388672 sec
W0605 17:38:50.230334 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004317760467529297 sec
W0605 17:38:50.231549 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00034928321838378906 sec
W0605 17:38:50.232604 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00027489662170410156 sec
W0605 17:38:50.234390 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00034546852111816406 sec
W0605 17:38:50.235389 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002620220184326172 sec
W0605 17:38:50.236262 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00040340423583984375 sec
W0605 17:38:50.237334 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.0004417896270751953 sec
W0605 17:38:50.238777 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00045418739318847656 sec
W0605 17:38:50.239361 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.001392364501953125 sec
W0605 17:38:50.240437 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.00044536590576171875 sec
I0605 17:38:50.241358 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.242370 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0005207061767578125 sec
I0605 17:38:50.243339 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.245457 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0005881786346435547 sec
W0605 17:38:50.246164 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0016107559204101562 sec
W0605 17:38:50.247194 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00035190582275390625 sec
W0605 17:38:50.249184 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00042319297790527344 sec
W0605 17:38:50.250619 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004248619079589844 sec
W0605 17:38:50.251869 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0006654262542724609 sec
W0605 17:38:50.253045 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0005459785461425781 sec
I0605 17:38:50.265692 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/combined_qkv/w with shape=[3, 2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
W0605 17:38:50.268131 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0005102157592773438 sec
W0605 17:38:50.269049 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003743171691894531 sec
W0605 17:38:50.269925 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00040221214294433594 sec
W0605 17:38:50.270813 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003933906555175781 sec
W0605 17:38:50.271438 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.004752159118652344 sec
W0605 17:38:50.272136 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.005776882171630859 sec
W0605 17:38:50.272483 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.006403684616088867 sec
W0605 17:38:50.273355 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00045037269592285156 sec
W0605 17:38:50.276827 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0007164478302001953 sec
I0605 17:38:50.278644 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/per_dim_scale/per_dim_scale with shape=[64], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.279554 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00043129920959472656 sec
W0605 17:38:50.282353 139756486557120 dispatch.py:272] Finished tracing + transforming logaddexp for pjit in 0.0012598037719726562 sec
W0605 17:38:50.282917 139756486557120 dispatch.py:272] Finished tracing + transforming softplus for pjit in 0.002315044403076172 sec
W0605 17:38:50.283740 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038123130798339844 sec
W0605 17:38:50.284803 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0006012916564941406 sec
W0605 17:38:50.286297 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.00058746337890625 sec
W0605 17:38:50.287535 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00041556358337402344 sec
W0605 17:38:50.288473 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038361549377441406 sec
W0605 17:38:50.290702 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0014166831970214844 sec
W0605 17:38:50.291424 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0024847984313964844 sec
W0605 17:38:50.292345 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0004284381866455078 sec
W0605 17:38:50.293292 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.0004420280456542969 sec
W0605 17:38:50.294705 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00046181678771972656 sec
W0605 17:38:50.295297 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.00140380859375 sec
W0605 17:38:50.296829 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00034689903259277344 sec
W0605 17:38:50.297558 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002791881561279297 sec
W0605 17:38:50.298322 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003561973571777344 sec
W0605 17:38:50.300544 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0005462169647216797 sec
W0605 17:38:50.301525 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003635883331298828 sec
W0605 17:38:50.302242 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002760887145996094 sec
W0605 17:38:50.303315 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0006041526794433594 sec
W0605 17:38:50.304219 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003628730773925781 sec
W0605 17:38:50.305894 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006916522979736328 sec
I0605 17:38:50.306929 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/post/w with shape=[2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
W0605 17:38:50.309180 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003814697265625 sec
W0605 17:38:50.310077 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00038695335388183594 sec
W0605 17:38:50.310939 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003902912139892578 sec
W0605 17:38:50.311874 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004658699035644531 sec
W0605 17:38:50.312483 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.004551410675048828 sec
W0605 17:38:50.313171 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.005565643310546875 sec
W0605 17:38:50.313507 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.006251096725463867 sec
W0605 17:38:50.314344 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004253387451171875 sec
W0605 17:38:50.317789 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006120204925537109 sec
W0605 17:38:50.321920 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00035834312438964844 sec
W0605 17:38:50.322905 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004630088806152344 sec
W0605 17:38:50.325140 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004410743713378906 sec
W0605 17:38:50.326235 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00044727325439453125 sec
W0605 17:38:50.328002 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.0004172325134277344 sec
W0605 17:38:50.328908 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003478527069091797 sec
W0605 17:38:50.330873 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.000431060791015625 sec
W0605 17:38:50.331428 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.001316070556640625 sec
I0605 17:38:50.345351 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.346637 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.359563 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/linear/w with shape=[2048, 8192], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
W0605 17:38:50.361818 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00038051605224609375 sec
W0605 17:38:50.363132 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00048160552978515625 sec
W0605 17:38:50.363966 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003731250762939453 sec
W0605 17:38:50.364569 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.0040357112884521484 sec
W0605 17:38:50.365290 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.005079507827758789 sec
W0605 17:38:50.365629 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.005695819854736328 sec
W0605 17:38:50.366461 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004315376281738281 sec
W0605 17:38:50.369798 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006268024444580078 sec
I0605 17:38:50.370664 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/bias/b with shape=[8192], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.371565 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00043845176696777344 sec
W0605 17:38:50.373037 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0005903244018554688 sec
W0605 17:38:50.374403 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004355907440185547 sec
W0605 17:38:50.375318 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003592967987060547 sec
W0605 17:38:50.376131 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003597736358642578 sec
W0605 17:38:50.376911 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002892017364501953 sec
W0605 17:38:50.377897 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0005381107330322266 sec
W0605 17:38:50.379199 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003714561462402344 sec
W0605 17:38:50.380480 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003561973571777344 sec
W0605 17:38:50.381428 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.00043392181396484375 sec
W0605 17:38:50.382789 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00045108795166015625 sec
W0605 17:38:50.383347 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0013604164123535156 sec
I0605 17:38:50.391109 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/linear/w with shape=[8192, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
W0605 17:38:50.393342 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00039124488830566406 sec
W0605 17:38:50.394674 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004923343658447266 sec
W0605 17:38:50.395514 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00037169456481933594 sec
W0605 17:38:50.396109 139756486557120 dispatch.py:272] Finished tracing + transforming _uniform for pjit in 0.004067897796630859 sec
W0605 17:38:50.396795 139756486557120 dispatch.py:272] Finished tracing + transforming _normal_real for pjit in 0.005061626434326172 sec
W0605 17:38:50.397140 139756486557120 dispatch.py:272] Finished tracing + transforming _normal for pjit in 0.0056803226470947266 sec
W0605 17:38:50.397976 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004203319549560547 sec
W0605 17:38:50.401306 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006275177001953125 sec
I0605 17:38:50.402132 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/bias/b with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.458132 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.459451 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.476639 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/combined_qkv/w with shape=[3, 2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:38:50.481575 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/per_dim_scale/per_dim_scale with shape=[64], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.493132 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/self_attention/post/w with shape=[2048, 32, 64], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:38:50.520059 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.521409 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/layer_norm/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.535170 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/linear/w with shape=[2048, 8192], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:38:50.539281 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer1/bias/b with shape=[8192], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.553280 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/linear/w with shape=[8192, 2048], dtype=<class 'jax.numpy.float32'>, init method=gaussian and scale=0.023
I0605 17:38:50.557278 139756486557120 base_layer.py:632] Creating var /lm/transformer/repeat/remat(scan(map_variables(map_variables(map_variables(sub)))))/x_layers_0/ff_layer/ffn_layer2/bias/b with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.586362 139756486557120 dispatch.py:272] Finished tracing + transforming logaddexp for pjit in 0.0009326934814453125 sec
W0605 17:38:50.587100 139756486557120 dispatch.py:272] Finished tracing + transforming real for pjit in 0.00015854835510253906 sec
W0605 17:38:50.588630 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002460479736328125 sec
W0605 17:38:50.589258 139756486557120 dispatch.py:272] Finished tracing + transforming real for pjit in 0.00016546249389648438 sec
I0605 17:38:50.685316 139756486557120 base_layer.py:632] Creating var /lm/final_ln/scale with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
I0605 17:38:50.686667 139756486557120 base_layer.py:632] Creating var /lm/final_ln/bias with shape=[2048], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.707986 139756486557120 dispatch.py:272] Finished tracing + transforming _einsum for pjit in 0.0006191730499267578 sec
I0605 17:38:50.711300 139756486557120 base_layer.py:632] Creating var /lm/softmax/logits_ffn/bias/b with shape=[51200], dtype=<class 'jax.numpy.float32'>, init method=constant and scale=0.0
W0605 17:38:50.712251 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00042057037353515625 sec
W0605 17:38:50.713695 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0005605220794677734 sec
W0605 17:38:50.716571 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00040793418884277344 sec
W0605 17:38:50.718845 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00027251243591308594 sec
W0605 17:38:50.721838 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004107952117919922 sec
W0605 17:38:50.724678 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0005333423614501953 sec
W0605 17:38:50.725696 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00043582916259765625 sec
W0605 17:38:50.726395 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002675056457519531 sec
W0605 17:38:50.727391 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0005466938018798828 sec
W0605 17:38:50.728161 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00026035308837890625 sec
W0605 17:38:50.728855 139756486557120 dispatch.py:272] Finished tracing + transforming log_softmax for pjit in 0.004996776580810547 sec
W0605 17:38:50.734375 139756486557120 dispatch.py:272] Finished tracing + transforming _squeeze for pjit in 0.0002372264862060547 sec
W0605 17:38:50.735692 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00036454200744628906 sec
W0605 17:38:50.736228 139756486557120 dispatch.py:272] Finished tracing + transforming _one_hot for pjit in 0.0014050006866455078 sec
W0605 17:38:50.737056 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003540515899658203 sec
W0605 17:38:50.739387 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00043201446533203125 sec
W0605 17:38:50.741913 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00024437904357910156 sec
W0605 17:38:50.743786 139756486557120 dispatch.py:272] Finished tracing + transforming _argmax for pjit in 0.00027298927307128906 sec
W0605 17:38:50.745871 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036263465881347656 sec
W0605 17:38:50.748253 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042724609375 sec
W0605 17:38:50.750594 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00039267539978027344 sec
W0605 17:38:50.755076 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00043892860412597656 sec
W0605 17:38:50.757273 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003552436828613281 sec
W0605 17:38:50.759236 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003609657287597656 sec
W0605 17:38:50.760184 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004966259002685547 sec
W0605 17:38:50.761510 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00043082237243652344 sec
W0605 17:38:50.762891 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00024199485778808594 sec
W0605 17:38:50.767537 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003418922424316406 sec
W0605 17:38:50.782197 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0004863739013671875 sec
W0605 17:38:50.786458 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0004134178161621094 sec
W0605 17:38:50.908352 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004239082336425781 sec
W0605 17:38:50.910371 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00039196014404296875 sec
W0605 17:38:50.911374 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004189014434814453 sec
W0605 17:38:50.912329 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00042819976806640625 sec
W0605 17:38:50.913593 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003299713134765625 sec
W0605 17:38:50.913979 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0010998249053955078 sec
W0605 17:38:50.915056 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004069805145263672 sec
W0605 17:38:50.915954 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003898143768310547 sec
W0605 17:38:50.918364 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004146099090576172 sec
W0605 17:38:50.920581 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0006053447723388672 sec
W0605 17:38:50.921390 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003533363342285156 sec
W0605 17:38:50.922576 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003867149353027344 sec
W0605 17:38:50.924040 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0005407333374023438 sec
W0605 17:38:50.925589 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033211708068847656 sec
W0605 17:38:50.926552 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00044608116149902344 sec
W0605 17:38:50.927937 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034737586975097656 sec
W0605 17:38:50.928853 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004279613494873047 sec
W0605 17:38:50.929699 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004150867462158203 sec
W0605 17:38:50.930606 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004134178161621094 sec
W0605 17:38:50.931371 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003495216369628906 sec
W0605 17:38:50.932270 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00041294097900390625 sec
W0605 17:38:50.933025 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034499168395996094 sec
W0605 17:38:50.933944 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004177093505859375 sec
W0605 17:38:50.934698 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034165382385253906 sec
W0605 17:38:50.935618 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004241466522216797 sec
W0605 17:38:50.936379 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003333091735839844 sec
W0605 17:38:50.937386 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0005240440368652344 sec
W0605 17:38:50.938133 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003476142883300781 sec
W0605 17:38:50.939027 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004146099090576172 sec
W0605 17:38:50.942208 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00035643577575683594 sec
W0605 17:38:50.943121 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042128562927246094 sec
W0605 17:38:50.943864 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003342628479003906 sec
W0605 17:38:50.944764 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042128562927246094 sec
W0605 17:38:50.946032 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0008378028869628906 sec
W0605 17:38:50.946952 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042057037353515625 sec
W0605 17:38:50.949638 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00041675567626953125 sec
W0605 17:38:50.950620 139756486557120 dispatch.py:272] Finished tracing + transforming isfinite for pjit in 0.00022411346435546875 sec
W0605 17:38:50.951457 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_all for pjit in 0.0004432201385498047 sec
W0605 17:38:50.952910 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004074573516845703 sec
W0605 17:38:50.953750 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00035572052001953125 sec
W0605 17:38:50.955095 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003368854522705078 sec
W0605 17:38:50.955853 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00032639503479003906 sec
W0605 17:38:50.956612 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003349781036376953 sec
W0605 17:38:50.957367 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033402442932128906 sec
W0605 17:38:50.958131 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034165382385253906 sec
W0605 17:38:50.958882 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003368854522705078 sec
W0605 17:38:50.960838 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003952980041503906 sec
W0605 17:38:50.961642 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003399848937988281 sec
W0605 17:38:50.962403 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003294944763183594 sec
W0605 17:38:50.974797 139756486557120 optimizers.py:1170] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
I0605 17:38:50.974858 139756486557120 optimizers.py:1173] Using sharded_adam.
W0605 17:38:50.974896 139756486557120 optimizers.py:580] DEPRECATION WARNING: p.weight_decay will be deprecated. In future, we will do a migration to remove p.weight_decay and after that, setting it will throw an exception. In future, we will use p.l2_regularizer_weight for coupled weight decay (i.e., weight decays that affect optimizer slots), and use p.decoupled_weight_decay for decoupled weight decay (i.e., weight decays that are added only to the final update).
W0605 17:38:50.976279 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003859996795654297 sec
W0605 17:38:50.977996 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.0002586841583251953 sec
W0605 17:38:50.979636 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0008521080017089844 sec
W0605 17:38:50.980117 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0016627311706542969 sec
W0605 17:38:50.981424 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.004053592681884766 sec
W0605 17:38:50.982840 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.0002465248107910156 sec
W0605 17:38:50.984083 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00047087669372558594 sec
W0605 17:38:50.984557 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012707710266113281 sec
W0605 17:38:50.985852 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0035409927368164062 sec
W0605 17:38:50.986913 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00024271011352539062 sec
W0605 17:38:50.988064 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003974437713623047 sec
W0605 17:38:50.988533 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011773109436035156 sec
W0605 17:38:50.989825 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0034170150756835938 sec
W0605 17:38:50.990873 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00024271011352539062 sec
W0605 17:38:50.992126 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003910064697265625 sec
W0605 17:38:50.992585 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012733936309814453 sec
W0605 17:38:50.993876 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0034990310668945312 sec
W0605 17:38:50.995401 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00038313865661621094 sec
W0605 17:38:50.996306 139756486557120 dispatch.py:272] Finished tracing + transforming _power for pjit in 0.00038886070251464844 sec
W0605 17:38:50.997263 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00040459632873535156 sec
W0605 17:38:51.002679 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003349781036376953 sec
W0605 17:38:51.003729 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003273487091064453 sec
W0605 17:38:51.019937 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034332275390625 sec
W0605 17:38:51.021034 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036072731018066406 sec
W0605 17:38:51.029161 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003571510314941406 sec
W0605 17:38:51.030206 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003421306610107422 sec
W0605 17:38:51.038254 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034236907958984375 sec
W0605 17:38:51.039307 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003333091735839844 sec
W0605 17:38:51.041758 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.000400543212890625 sec
W0605 17:38:51.042524 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002651214599609375 sec
W0605 17:38:51.044231 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003452301025390625 sec
W0605 17:38:51.046305 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038170814514160156 sec
W0605 17:38:51.047033 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002536773681640625 sec
W0605 17:38:51.048109 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00033783912658691406 sec
W0605 17:38:51.048935 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038361549377441406 sec
W0605 17:38:51.049658 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002503395080566406 sec
W0605 17:38:51.050806 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00040531158447265625 sec
W0605 17:38:51.051653 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003936290740966797 sec
W0605 17:38:51.052373 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00025343894958496094 sec
W0605 17:38:51.053460 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003292560577392578 sec
W0605 17:38:51.054191 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.0002491474151611328 sec
W0605 17:38:51.055417 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004432201385498047 sec
W0605 17:38:51.055958 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.001308441162109375 sec
W0605 17:38:51.057440 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.0002446174621582031 sec
W0605 17:38:51.059051 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0022339820861816406 sec
W0605 17:38:51.060664 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003437995910644531 sec
W0605 17:38:51.063558 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00025272369384765625 sec
W0605 17:38:51.064804 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00048232078552246094 sec
W0605 17:38:51.065351 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0013427734375 sec
W0605 17:38:51.067524 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003342628479003906 sec
W0605 17:38:51.068215 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.0002677440643310547 sec
W0605 17:38:51.069389 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00043010711669921875 sec
W0605 17:38:51.069930 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0012748241424560547 sec
W0605 17:38:51.072082 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00033545494079589844 sec
W0605 17:38:51.072823 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00032258033752441406 sec
W0605 17:38:51.073729 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004303455352783203 sec
W0605 17:38:51.076472 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003464221954345703 sec
W0605 17:38:51.077437 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00040221214294433594 sec
W0605 17:38:51.080010 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003974437713623047 sec
W0605 17:38:51.084977 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003871917724609375 sec
W0605 17:38:51.085875 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003821849822998047 sec
W0605 17:38:51.106950 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00027370452880859375 sec
W0605 17:38:51.108178 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0004146099090576172 sec
W0605 17:38:51.108645 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.001222848892211914 sec
W0605 17:38:51.109961 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0036516189575195312 sec
W0605 17:38:51.113869 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00024700164794921875 sec
W0605 17:38:51.115013 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003838539123535156 sec
W0605 17:38:51.115498 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.00118255615234375 sec
W0605 17:38:51.116794 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.0034394264221191406 sec
W0605 17:38:51.123821 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00025463104248046875 sec
W0605 17:38:51.125084 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0004799365997314453 sec
W0605 17:38:51.125572 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012826919555664062 sec
W0605 17:38:51.126914 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.003610372543334961 sec
W0605 17:38:51.132493 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.00025391578674316406 sec
W0605 17:38:51.134341 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0010538101196289062 sec
W0605 17:38:51.134835 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0018856525421142578 sec
W0605 17:38:51.136130 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.004152536392211914 sec
W0605 17:38:51.140041 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.0002532005310058594 sec
W0605 17:38:51.141291 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0004820823669433594 sec
W0605 17:38:51.141758 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012583732604980469 sec
W0605 17:38:51.143049 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.003523588180541992 sec
W0605 17:38:51.146885 139756486557120 dispatch.py:272] Finished tracing + transforming isnan for pjit in 0.0002493858337402344 sec
W0605 17:38:51.148137 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0004699230194091797 sec
W0605 17:38:51.148612 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012650489807128906 sec
W0605 17:38:51.149911 139756486557120 dispatch.py:272] Finished tracing + transforming nan_to_num for pjit in 0.003537416458129883 sec
W0605 17:38:51.164990 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003509521484375 sec
W0605 17:38:51.167197 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033211708068847656 sec
W0605 17:38:51.179063 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034117698669433594 sec
W0605 17:38:51.181202 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036454200744628906 sec
W0605 17:38:51.205911 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00033593177795410156 sec
W0605 17:38:51.208079 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003936290740966797 sec
W0605 17:38:51.267738 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00036597251892089844 sec
W0605 17:38:51.269978 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003571510314941406 sec
W0605 17:38:51.283516 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003502368927001953 sec
W0605 17:38:51.295324 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003483295440673828 sec
W0605 17:38:51.297498 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034880638122558594 sec
W0605 17:38:51.301339 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00040984153747558594 sec
W0605 17:38:51.302777 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002617835998535156 sec
W0605 17:38:51.304450 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00033354759216308594 sec
W0605 17:38:51.305875 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003902912139892578 sec
W0605 17:38:51.307199 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00026226043701171875 sec
W0605 17:38:51.308934 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0004074573516845703 sec
W0605 17:38:51.313358 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00040531158447265625 sec
W0605 17:38:51.314678 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002655982971191406 sec
W0605 17:38:51.316417 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003962516784667969 sec
W0605 17:38:51.324674 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003845691680908203 sec
W0605 17:38:51.326095 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003135204315185547 sec
W0605 17:38:51.327770 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.000331878662109375 sec
W0605 17:38:51.329179 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003993511199951172 sec
W0605 17:38:51.330489 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002593994140625 sec
W0605 17:38:51.332172 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.000331878662109375 sec
W0605 17:38:51.333650 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038313865661621094 sec
W0605 17:38:51.334960 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002560615539550781 sec
W0605 17:38:51.336631 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00034332275390625 sec
W0605 17:38:51.337970 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.000255584716796875 sec
W0605 17:38:51.339765 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0005276203155517578 sec
W0605 17:38:51.340311 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0014069080352783203 sec
W0605 17:38:51.347506 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003342628479003906 sec
W0605 17:38:51.349030 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00026154518127441406 sec
W0605 17:38:51.350681 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004153251647949219 sec
W0605 17:38:51.351207 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.001252889633178711 sec
W0605 17:38:51.354593 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003325939178466797 sec
W0605 17:38:51.360800 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00024509429931640625 sec
W0605 17:38:51.362443 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042128562927246094 sec
W0605 17:38:51.362977 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0012655258178710938 sec
W0605 17:38:51.366878 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003380775451660156 sec
W0605 17:38:51.381033 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.0002589225769042969 sec
W0605 17:38:51.382707 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00042128562927246094 sec
W0605 17:38:51.383244 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0012717247009277344 sec
W0605 17:38:51.386665 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0004096031188964844 sec
W0605 17:38:51.388154 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.0002455711364746094 sec
W0605 17:38:51.389823 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.0004107952117919922 sec
W0605 17:38:51.390361 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0012784004211425781 sec
W0605 17:38:51.393694 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003459453582763672 sec
W0605 17:38:51.395206 139756486557120 dispatch.py:272] Finished tracing + transforming square for pjit in 0.00025153160095214844 sec
W0605 17:38:51.396828 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_sum for pjit in 0.00041174888610839844 sec
W0605 17:38:51.397366 139756486557120 dispatch.py:272] Finished tracing + transforming _mean for pjit in 0.0012562274932861328 sec
W0605 17:38:51.400723 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00036334991455078125 sec
W0605 17:38:51.402534 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00038313865661621094 sec
W0605 17:38:51.413817 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003960132598876953 sec
W0605 17:38:51.458365 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00040435791015625 sec
W0605 17:38:51.458882 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012650489807128906 sec
W0605 17:38:51.460613 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003859996795654297 sec
W0605 17:38:51.461089 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.001180410385131836 sec
W0605 17:38:51.462345 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003674030303955078 sec
W0605 17:38:51.462806 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011293888092041016 sec
W0605 17:38:51.464207 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003802776336669922 sec
W0605 17:38:51.464680 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012753009796142578 sec
W0605 17:38:51.465978 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003809928894042969 sec
W0605 17:38:51.466452 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011639595031738281 sec
W0605 17:38:51.467733 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00038886070251464844 sec
W0605 17:38:51.468206 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011620521545410156 sec
W0605 17:38:51.469493 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0004017353057861328 sec
W0605 17:38:51.469960 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011703968048095703 sec
W0605 17:38:51.471318 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003769397735595703 sec
W0605 17:38:51.471788 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012331008911132812 sec
W0605 17:38:51.474819 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003876686096191406 sec
W0605 17:38:51.475300 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011796951293945312 sec
W0605 17:38:51.476653 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00045108795166015625 sec
W0605 17:38:51.477125 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0012183189392089844 sec
W0605 17:38:51.478438 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.000385284423828125 sec
W0605 17:38:51.478904 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011568069458007812 sec
W0605 17:38:51.480214 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.00021004676818847656 sec
W0605 17:38:51.480560 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0008604526519775391 sec
W0605 17:38:51.484501 139756486557120 dispatch.py:272] Finished tracing + transforming _broadcast_arrays for pjit in 0.0003781318664550781 sec
W0605 17:38:51.484953 139756486557120 dispatch.py:272] Finished tracing + transforming _where for pjit in 0.0011394023895263672 sec
W0605 17:38:51.518642 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00034880638122558594 sec
W0605 17:38:51.519460 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003371238708496094 sec
W0605 17:38:51.520232 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003368854522705078 sec
W0605 17:38:51.520990 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003304481506347656 sec
W0605 17:38:51.522938 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004017353057861328 sec
W0605 17:38:51.523720 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003342628479003906 sec
W0605 17:38:51.524483 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003447532653808594 sec
W0605 17:38:51.526116 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004146099090576172 sec
W0605 17:38:51.526935 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002446174621582031 sec
W0605 17:38:51.527758 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0003864765167236328 sec
W0605 17:38:51.528560 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003807544708251953 sec
W0605 17:38:51.533638 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002372264862060547 sec
W0605 17:38:51.534479 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00039267539978027344 sec
W0605 17:38:51.537130 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002434253692626953 sec
W0605 17:38:51.538020 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00043773651123046875 sec
W0605 17:38:51.540642 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002391338348388672 sec
W0605 17:38:51.541481 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00040149688720703125 sec
W0605 17:38:51.544121 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.00024962425231933594 sec
W0605 17:38:51.545009 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00044608116149902344 sec
W0605 17:38:51.546122 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00033855438232421875 sec
W0605 17:38:51.548141 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.00024080276489257812 sec
W0605 17:38:51.548953 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00038123130798339844 sec
W0605 17:38:51.550048 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003368854522705078 sec
W0605 17:38:51.552147 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002999305725097656 sec
W0605 17:38:51.552977 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0003833770751953125 sec
W0605 17:38:51.554072 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.0003323554992675781 sec
W0605 17:38:51.556094 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.000232696533203125 sec
W0605 17:38:51.556920 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0003819465637207031 sec
W0605 17:38:51.558031 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00033664703369140625 sec
W0605 17:38:51.570347 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.00024271011352539062 sec
W0605 17:38:51.571178 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.0003898143768310547 sec
W0605 17:38:51.572269 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00032901763916015625 sec
W0605 17:38:51.574327 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002353191375732422 sec
W0605 17:38:51.575155 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00037741661071777344 sec
W0605 17:38:51.576297 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00038695335388183594 sec
W0605 17:38:51.578324 139756486557120 dispatch.py:272] Finished tracing + transforming absolute for pjit in 0.0002453327178955078 sec
W0605 17:38:51.579146 139756486557120 dispatch.py:272] Finished tracing + transforming _reduce_max for pjit in 0.00038051605224609375 sec
W0605 17:38:51.580233 139756486557120 dispatch.py:272] Finished tracing + transforming true_divide for pjit in 0.00032901763916015625 sec
W0605 17:38:51.625336 139756486557120 dispatch.py:272] Finished tracing + transforming _wrapped_step_fn for pmap in 1.566312313079834 sec
W0605 17:38:51.626110 139756486557120 pxla.py:859] Compiling _wrapped_step_fn (139741169123712) for 1 devices with args (ShapedArray(uint32[1]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048,2048]), ShapedArray(float32[1,51200]), ShapedArray(float32[1,2048,51200]), ShapedArray(float32[1,24,8192]), ShapedArray(float32[1,24,2048,8192]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,8192,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,3,2048,32,64]), ShapedArray(float32[1,24,64]), ShapedArray(float32[1,24,2048,32,64]), ShapedArray(int32[1]), ShapedArray(int32[1]), ShapedArray(int32[1]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048,2048]), ShapedArray(float32[1,51200]), ShapedArray(float32[1,2048,51200]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048]), ShapedArray(float32[1,2048,2048]), ShapedArray(float32[1,51200]), ShapedArray(float32[1,2048,51200]), ShapedArray(int32[1]), ShapedArray(int32[1,24]), ShapedArray(int32[1,24]), ShapedArray(int32[1,24]), ShapedArray(float32[1,24,8192]), ShapedArray(float32[1,24,2048,8192]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,8192,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,3,2048,32,64]), ShapedArray(float32[1,24,64]), ShapedArray(float32[1,24,2048,32,64]), ShapedArray(float32[1,24,8192]), ShapedArray(float32[1,24,2048,8192]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,8192,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,2048]), ShapedArray(float32[1,24,3,2048,32,64]), ShapedArray(float32[1,24,64]), ShapedArray(float32[1,24,2048,32,64]), ShapedArray(int32[1,24]), ShapedArray(uint32[1,2]), ShapedArray(float32[1,1]), ShapedArray(int32[1,1,2048]), ShapedArray(int32[1,1,2048]), ShapedArray(float32[1,1,2048]), ShapedArray(int32[1,1,2048]), ShapedArray(int32[1,1,2048]), ShapedArray(float32[1,1,2048])). (num_replicas=1)
/workspace/jax/jax/_src/interpreters/mlir.py:618: UserWarning: Some donated buffers were not usable: ShapedArray(uint32[]), ShapedArray(float32[2048]), ShapedArray(float32[2048]), ShapedArray(float32[2048,2048]), ShapedArray(float32[51200]), ShapedArray(float32[2048,51200]), ShapedArray(float32[24,8192]), ShapedArray(float32[24,2048,8192]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,8192,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,3,2048,32,64]), ShapedArray(float32[24,64]), ShapedArray(float32[24,2048,32,64]), ShapedArray(int32[]), ShapedArray(int32[]), ShapedArray(int32[]), ShapedArray(float32[2048]), ShapedArray(float32[2048]), ShapedArray(float32[2048,2048]), ShapedArray(float32[51200]), ShapedArray(float32[2048,51200]), ShapedArray(float32[2048]), ShapedArray(float32[2048]), ShapedArray(float32[2048,2048]), ShapedArray(float32[51200]), ShapedArray(float32[2048,51200]), ShapedArray(int32[]), ShapedArray(int32[24]), ShapedArray(int32[24]), ShapedArray(int32[24]), ShapedArray(float32[24,8192]), ShapedArray(float32[24,2048,8192]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,8192,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,3,2048,32,64]), ShapedArray(float32[24,64]), ShapedArray(float32[24,2048,32,64]), ShapedArray(float32[24,8192]), ShapedArray(float32[24,2048,8192]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,8192,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,2048]), ShapedArray(float32[24,3,2048,32,64]), ShapedArray(float32[24,64]), ShapedArray(float32[24,2048,32,64]), ShapedArray(int32[24]).
Donation is not implemented for iree_cuda.
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer-donation.
  warnings.warn(f"Some donated buffers were not usable: {', '.join(unused_donations)}.\n{msg}")
W0605 17:38:51.637015 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00033736228942871094 sec
W0605 17:38:51.637705 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_seed for pjit in 0.0017371177673339844 sec
W0605 17:38:51.639181 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.0001888275146484375 sec
W0605 17:38:51.639944 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0015056133270263672 sec
W0605 17:38:51.640912 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_fold_in for pjit in 0.0051746368408203125 sec
W0605 17:38:51.645074 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0004105567932128906 sec
W0605 17:38:51.646182 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003085136413574219 sec
W0605 17:38:51.647176 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003046989440917969 sec
W0605 17:38:51.648066 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0002963542938232422 sec
W0605 17:38:51.648782 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00030303001403808594 sec
W0605 17:38:51.700799 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00019788742065429688 sec
W0605 17:38:51.701670 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0016200542449951172 sec
W0605 17:38:51.702680 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.0029451847076416016 sec
W0605 17:38:51.706055 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003178119659423828 sec
W0605 17:38:51.707761 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0010182857513427734 sec
W0605 17:38:51.708686 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003020763397216797 sec
W0605 17:38:51.709419 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003161430358886719 sec
W0605 17:38:51.762471 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.00020575523376464844 sec
W0605 17:38:51.763350 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.0016407966613769531 sec
W0605 17:38:51.764356 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.002961874008178711 sec
W0605 17:38:51.767745 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003142356872558594 sec
W0605 17:38:51.768834 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0004010200500488281 sec
W0605 17:38:51.769746 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003020763397216797 sec
W0605 17:38:51.770464 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003070831298828125 sec
W0605 17:38:51.834810 139756486557120 dispatch.py:272] Finished tracing + transforming ravel for pjit in 0.0002052783966064453 sec
W0605 17:38:51.835697 139756486557120 dispatch.py:272] Finished tracing + transforming threefry_2x32 for pjit in 0.001672983169555664 sec
W0605 17:38:51.836712 139756486557120 dispatch.py:272] Finished tracing + transforming _threefry_split_original for pjit in 0.003003358840942383 sec
W0605 17:38:51.840270 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00044655799865722656 sec
W0605 17:38:51.841293 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003247261047363281 sec
W0605 17:38:51.842204 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.00030112266540527344 sec
W0605 17:38:51.842923 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003039836883544922 sec
W0605 17:38:51.904477 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003521442413330078 sec
W0605 17:38:51.906139 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003561973571777344 sec
W0605 17:38:51.930905 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.000316619873046875 sec
W0605 17:38:52.100980 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003898143768310547 sec
W0605 17:38:52.102146 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003616809844970703 sec
W0605 17:38:52.102983 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.00030803680419921875 sec
W0605 17:38:52.103872 139756486557120 dispatch.py:272] Finished tracing + transforming fn for pjit in 0.0003063678741455078 sec
W0605 17:38:52.131105 139756486557120 dispatch.py:272] Finished tracing + transforming <lambda> for pjit in 0.0003311634063720703 sec
W0605 17:38:52.923904 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion pmap(_wrapped_step_fn) in 1.2972416877746582 sec
W0605 17:39:41.445357 139756486557120 dispatch.py:272] Finished XLA compilation of _wrapped_step_fn in 48.50671124458313 sec
W0605 17:39:41.464271 139756486557120 dispatch.py:272] Finished tracing + transforming _multi_slice for pjit in 0.0005140304565429688 sec
W0605 17:39:41.465060 139756486557120 pxla.py:1882] Compiling _multi_slice for with global shapes and types [ShapedArray(uint32[1,2])]. Argument mapping: (GSPMDSharding({replicated}),).
W0605 17:39:41.470349 139756486557120 dispatch.py:272] Finished jaxpr to MLIR module conversion jit(_multi_slice) in 0.005116462707519531 sec
W0605 17:39:41.740634 139756486557120 dispatch.py:272] Finished XLA compilation of jit(_multi_slice) in 0.2699155807495117 sec
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600142a3c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600142a3c0, semaphore=0x6060001683e0, value=89 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=90, fence=0x6040003e4450 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600142a3c0, from_fence=0x606000359e40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008a3080, semaphore=0x6060001683e0, value=90 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000476ba0, f=0, wait_fence=0x60600142a3c0 {0x6060001683e0:89}, signal_fence=0x6040003e4450 {0x6060001683e0:90} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600142aa80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600142ab40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600142aa80, semaphore=0x6060001683e0, value=90 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600142ab40, semaphore=0x6060001683e0, value=90 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=4, buffer=0x60c003e70dc0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=4, wait={0x6060001682c0:132}, signal={0x6060001682c0:133} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:133}, signal={0x6060001682c0:134} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606002537d80 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606002537cc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606002537d80, semaphore=0x6060001682c0, value=134 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606002537cc0, semaphore=0x6060001682c0, value=134 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e70ac0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:134}, signal={0x6060001682c0:135} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:135}, signal={0x6060001682c0:136} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000280f40 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000542240 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000280f40, semaphore=0x6060001682c0, value=136 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000542240, semaphore=0x6060001682c0, value=136 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e707c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:136}, signal={0x6060001682c0:137} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:137}, signal={0x6060001682c0:138} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606001b12fe0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600072f0e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606001b12fe0, semaphore=0x6060001682c0, value=138 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072f0e0, semaphore=0x6060001682c0, value=138 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e704c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:138}, signal={0x6060001682c0:139} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:139}, signal={0x6060001682c0:140} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060014ae2a0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606001773980 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060014ae2a0, semaphore=0x6060001682c0, value=140 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606001773980, semaphore=0x6060001682c0, value=140 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e701c0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:140}, signal={0x6060001682c0:141} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:141}, signal={0x6060001682c0:142} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606001061900 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000bd0620 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606001061900, semaphore=0x6060001682c0, value=142 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bd0620, semaphore=0x6060001682c0, value=142 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e0fec0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:142}, signal={0x6060001682c0:143} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:143}, signal={0x6060001682c0:144} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606001013480 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x6060014ae480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606001013480, semaphore=0x6060001682c0, value=144 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060014ae480, semaphore=0x6060001682c0, value=144 (OK)
:: IREE INVOKE (hal_allocator_allocate_buffer): allocator=0x60b000158cc0, size=8192, buffer=0x60c003e0fbc0 (OK)
:: IREE INVOKE (hal_device_queue_alloca): device=0x6110011f9800, size=8192, wait={0x6060001682c0:144}, signal={0x6060001682c0:145} (OK)
:: IREE INVOKE (hal_device_queue_execute): device=0x6110011f9800, wait={0x6060001682c0:145}, signal={0x6060001682c0:146} (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606002c0c3c0 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x606000c9f680 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606002c0c3c0, semaphore=0x6060001682c0, value=146 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c9f680, semaphore=0x6060001682c0, value=146 (OK)
:: IREE INVOKE (hal_fence_create): capacity=2, fence=0x60600089edc0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600089edc0, semaphore=0x6060001683e0, value=90 (OK)
:: IREE INVOKE (hal_fence_create_at): semaphore=0x6060001683e0, value=91, fence=0x60400144bb50 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000a17c20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004cfdc0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600079a8a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a900, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600079a960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600079a9c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600079aa20 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600075fbc0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000681620 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b00c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006f03e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006811a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060002b7720 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681ec0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000680de0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681e60, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b0600 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000682940, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b2040 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b25e0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000681b00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0540, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b2100 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006824c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060002b77e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b1c80, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b1440 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0180, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b1860 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000681680, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000682ca0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2ac0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000681a40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b0f00, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000bbe920 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600081a4c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000bbea40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000753b00, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000a93c40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a93820, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006b29a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006b2b20, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000970700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600034d420, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000967e80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000967e20, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000ca0c40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006442a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060005719a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007cf6a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060005b9700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060005ba060, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000946700 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009467c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060007e8240 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000761a80, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060003915e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060001c5e60, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600044a9e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600061e4a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000927e00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000315740, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000965a80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060006dbb00, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060008210c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600086c7e0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060007a8160 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000503c60, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600009c7a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000779d80, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060004af600 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003eb820, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600049fc40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060004e61a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060004fbaa0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003a06a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060003f5780 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600089cfc0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060009c0200 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c41900, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060009c08c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009e2940, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000bca3e0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bca5c0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600065b940 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000659900, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600023cde0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060009040a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060001619c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060002a2c00, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006d96a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600035b820, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600086a5c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000cbc960, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006fdf40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000452720, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000c07160 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060003f8480, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060009a7960 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000d16de0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060005cd560 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060007c6280, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060002c5ee0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060008c5340, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600038bd00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b5d360, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060006d6160 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000a045a0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060005686a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b56e20, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000a1fae0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000b57240, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000abce00 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000771f80, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x60600142aa80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600142ab40, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606002537d80 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606002537cc0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606000280f40 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000542240, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606001b12fe0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x60600072f0e0, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x6060014ae2a0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606001773980, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606001061900 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000bd0620, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606001013480 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x6060014ae480, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (hal_fence_extend): into_fence=0x60600089edc0, from_fence=0x606002c0c3c0 (OK)
:: IREE INVOKE (hal_fence_insert): fence=0x606000c9f680, semaphore=0x6060001683e0, value=91 (OK)
:: IREE INVOKE (vm_invoke[async]): context=0x60b000906510, f=0, wait_fence=0x60600089edc0 {0x6060001683e0:90, 0x6060001682c0:146}, signal_fence=0x60400144bb50 {0x6060001683e0:91}=================================================================
==12037==ERROR: AddressSanitizer: use-after-poison on address 0x62d00047e400 at pc 0x7f1b9823db9a bp 0x7ffdadc04c50 sp 0x7ffdadc04420
WRITE of size 72 at 0x62d00047e400 thread T0
    #0 0x7f1b9823db99 in __asan_memcpy (/usr/lib/llvm-14/lib/clang/14.0.0/lib/linux/libclang_rt.asan-x86_64.so+0xccb99) (BuildId: 0fc20d2022c0d572c45850b6d559d9ccfabd5443)
    #1 0x7f19603f099f in iree_hal_collective_batch_append /proc/self/cwd/external/iree_core/runtime/src/iree/hal/utils/collective_batch.c:102:36
    #2 0x7f19603d7aab in iree_hal_cuda_stream_command_buffer_collective /proc/self/cwd/external/iree_core/runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c:389:10
    #3 0x7f196041aa42 in iree_hal_command_buffer_collective /proc/self/cwd/external/iree_core/runtime/src/iree/hal/command_buffer.c:494:26
    #4 0x7f19603f3727 in iree_hal_deferred_command_buffer_apply_collective /proc/self/cwd/external/iree_core/runtime/src/iree/hal/utils/deferred_command_buffer.c:645:10
    #5 0x7f19603f18c0 in iree_hal_deferred_command_buffer_apply /proc/self/cwd/external/iree_core/runtime/src/iree/hal/utils/deferred_command_buffer.c:923:16
    #6 0x7f19603b01cd in iree_hal_cuda_device_queue_execute /proc/self/cwd/external/iree_core/runtime/src/iree/hal/drivers/cuda/cuda_device.c:527:7
    #7 0x7f19604254e4 in iree_hal_device_queue_execute /proc/self/cwd/external/iree_core/runtime/src/iree/hal/device.c:249:26
    #8 0x7f195e9d5dfe in iree_hal_module_device_queue_execute /proc/self/cwd/external/iree_core/runtime/src/iree/modules/hal/module.c:1014:10
    #9 0x7f195eb4fd95 in iree_vm_shim_rIrrCrD_v /proc/self/cwd/external/iree_core/runtime/src/iree/vm/shims.c:68:1
    #10 0x7f195eb31d97 in iree_vm_native_module_issue_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/native_module.c:338:7
    #11 0x7f195eb309e9 in iree_vm_native_module_begin_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/native_module.c:392:10
    #12 0x7f195ea4632c in iree_vm_bytecode_issue_import_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:452:7
    #13 0x7f195ea41d2e in iree_vm_bytecode_call_import_variadic /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:609:10
    #14 0x7f195ea29d9d in iree_vm_bytecode_dispatch /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:1667:5
    #15 0x7f195e9ff564 in iree_vm_bytecode_dispatch_begin /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:636:10
    #16 0x7f195e9f3fd0 in iree_vm_bytecode_module_begin_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/module.c:779:10
    #17 0x7f195eb10c2e in iree_vm_begin_invoke /proc/self/cwd/external/iree_core/runtime/src/iree/vm/invocation.c:504:7
    #18 0x7f195eb0e64a in iree_vm_invoke /proc/self/cwd/external/iree_core/runtime/src/iree/vm/invocation.c:302:26
    #19 0x7f195e95abb8 in iree::pjrt::LoadedExecutableInstance::BatchExecute(PJRT_LoadedExecutable_Execute_Args*) /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1797:9
    #20 0x7f195e964d8d in iree::pjrt::LoadedExecutableInstance::BindApi(PJRT_Api*)::$_54::operator()(PJRT_LoadedExecutable_Execute_Args*) const /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1590:61
    #21 0x7f195e964d34 in iree::pjrt::LoadedExecutableInstance::BindApi(PJRT_Api*)::$_54::__invoke(PJRT_LoadedExecutable_Execute_Args*) /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1587:8
    #22 0x7f1b8e11a00a in xla::PjRtCApiLoadedExecutable::Execute(absl::lts_20230125::Span<std::vector<xla::PjRtBuffer*, std::allocator<xla::PjRtBuffer*> > const>, xla::ExecuteOptions const&, std::optional<std::vector<xla::PjRtFuture<absl::lts_20230125::Status>, std::allocator<xla::PjRtFuture<absl::lts_20230125::Status> > > >&) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0xcaf00a) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #23 0x7f1b906deca6 in xla::ifrt::PjRtLoadedExecutable::Execute(absl::lts_20230125::Span<tsl::RCReference<xla::ifrt::Array> >, xla::ExecuteOptions const&, std::optional<xla::ifrt::DeviceList>) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x3273ca6) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #24 0x7f1b8e0c8836 in absl::lts_20230125::StatusOr<xla::PyExecuteResults> xla::(anonymous namespace)::ExecuteShardedOnLocalDevicesInternal<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, xla::(anonymous namespace)::ShardedBufferAdapter<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >(xla::ExecuteOptions const&, std::shared_ptr<xla::PyClient> const&, xla::ifrt::LoadedExecutable*, absl::lts_20230125::Span<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > const>, std::optional<std::vector<xla::PjRtFuture<absl::lts_20230125::Status>, std::allocator<xla::PjRtFuture<absl::lts_20230125::Status> > > >&) py_executable.cc
    #25 0x7f1b8e0c9b8d in xla::PyLoadedExecutable::ExecuteSharded(std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0xc5eb8d) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #26 0x7f1b8dd9b1d3 in void pybind11::cpp_function::initialize<xla::ValueOrThrowWrapper<absl::lts_20230125::StatusOr<xla::PyExecuteResults> (std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), xla::PyLoadedExecutable>, xla::PyExecuteResults, xla::PyLoadedExecutable&, std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v>(xla::ValueOrThrowWrapper<absl::lts_20230125::StatusOr<xla::PyExecuteResults> (std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), xla::PyLoadedExecutable>&&, xla::PyExecuteResults (*)(xla::PyLoadedExecutable&, std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&)::'lambda1'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x9301d3) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #27 0x7f1b8dd6fed0 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x904ed0) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #28 0x55e95b1b499d  (/usr/bin/python3.10+0x15c99d) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #29 0x55e95b1ab4aa in _PyObject_MakeTpCall (/usr/bin/python3.10+0x1534aa) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #30 0x55e95b1c2f0a  (/usr/bin/python3.10+0x16af0a) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #31 0x55e95b1a3461 in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x14b461) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #32 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #33 0x55e95b19faef in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147aef) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #34 0x55e95b1aa633 in _PyObject_FastCallDictTstate (/usr/bin/python3.10+0x152633) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #35 0x55e95b1bfd10 in _PyObject_Call_Prepend (/usr/bin/python3.10+0x167d10) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #36 0x55e95b2dd60f  (/usr/bin/python3.10+0x28560f) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #37 0x55e95b1c387a in PyObject_Call (/usr/bin/python3.10+0x16b87a) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #38 0x55e95b19faef in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147aef) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #39 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #40 0x55e95b19faef in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147aef) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #41 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #42 0x7f1b8de43a4c in jax::PmapFunction::Call(pybind11::handle, _object* const*, unsigned long, _object*) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x9d8a4c) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #43 0x7f1b8de4424a in JaxPmapFunction_tp_vectorcall (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x9d924a) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #44 0x55e95b19d784 in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x145784) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #45 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #46 0x55e95b19d784 in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x145784) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #47 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #48 0x55e95b19d8ca in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458ca) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #49 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #50 0x55e95b19d8ca in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458ca) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #51 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #52 0x55e95b1c38e1 in PyObject_Call (/usr/bin/python3.10+0x16b8e1) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #53 0x55e95b19faef in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x147aef) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #54 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #55 0x55e95b19d8ca in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x1458ca) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #56 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #57 0x55e95b19eade in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x146ade) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #58 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #59 0x55e95b19eade in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x146ade) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #60 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #61 0x55e95b19eade in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x146ade) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #62 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #63 0x55e95b19d784 in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x145784) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #64 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #65 0x55e95b19d784 in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x145784) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #66 0x55e95b1b51eb in _PyFunction_Vectorcall (/usr/bin/python3.10+0x15d1eb) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #67 0x55e95b19eade in _PyEval_EvalFrameDefault (/usr/bin/python3.10+0x146ade) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #68 0x55e95b199ed5  (/usr/bin/python3.10+0x141ed5) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #69 0x55e95b290365 in PyEval_EvalCode (/usr/bin/python3.10+0x238365) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #70 0x55e95b2bd107  (/usr/bin/python3.10+0x265107) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #71 0x55e95b2b5f5a  (/usr/bin/python3.10+0x25df5a) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #72 0x55e95b2bce54  (/usr/bin/python3.10+0x264e54) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #73 0x55e95b2bc337 in _PyRun_SimpleFileObject (/usr/bin/python3.10+0x264337) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #74 0x55e95b2bc032 in _PyRun_AnyFileObject (/usr/bin/python3.10+0x264032) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #75 0x55e95b2ad2dd in Py_RunMain (/usr/bin/python3.10+0x2552dd) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #76 0x55e95b28332c in Py_BytesMain (/usr/bin/python3.10+0x22b32c) (BuildId: 148e086667839ef13939196984d6f717c331bd76)
    #77 0x7f1b97e3ed8f  (/usr/lib/x86_64-linux-gnu/libc.so.6+0x29d8f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #78 0x7f1b97e3ee3f in __libc_start_main (/usr/lib/x86_64-linux-gnu/libc.so.6+0x29e3f) (BuildId: 69389d485a9793dbe873f0ea2c93e02efaa9aa3d)
    #79 0x55e95b283224 in _start (/usr/bin/python3.10+0x22b224) (BuildId: 148e086667839ef13939196984d6f717c331bd76)

0x62d00047e400 is located 0 bytes inside of 32768-byte region [0x62d00047e400,0x62d000486400)
allocated by thread T0 here:
    #0 0x7f1b9823e7ee in __interceptor_malloc (/usr/lib/llvm-14/lib/clang/14.0.0/lib/linux/libclang_rt.asan-x86_64.so+0xcd7ee) (BuildId: 0fc20d2022c0d572c45850b6d559d9ccfabd5443)
    #1 0x7f1960444f70 in iree_allocator_system_alloc /proc/self/cwd/external/iree_core/runtime/src/iree/base/allocator.c:88:17
    #2 0x7f1960444a39 in iree_allocator_system_ctl /proc/self/cwd/external/iree_core/runtime/src/iree/base/allocator.c:126:14
    #3 0x7f1960443ca8 in iree_allocator_issue_alloc /proc/self/cwd/external/iree_core/runtime/src/iree/base/allocator.c:27:10
    #4 0x7f1960443f44 in iree_allocator_malloc_uninitialized /proc/self/cwd/external/iree_core/runtime/src/iree/base/allocator.c:38:10
    #5 0x7f19603fc567 in iree_arena_block_pool_acquire /proc/self/cwd/external/iree_core/runtime/src/iree/base/internal/arena.c:77:5
    #6 0x7f19603fd455 in iree_arena_allocate /proc/self/cwd/external/iree_core/runtime/src/iree/base/internal/arena.c:187:5
    #7 0x7f19603f9cc6 in iree_hal_cmd_list_append_command /proc/self/cwd/external/iree_core/runtime/src/iree/hal/utils/deferred_command_buffer.c:113:3
    #8 0x7f19603f7eb3 in iree_hal_deferred_command_buffer_push_constants /proc/self/cwd/external/iree_core/runtime/src/iree/hal/utils/deferred_command_buffer.c:672:3
    #9 0x7f196041ad49 in iree_hal_command_buffer_push_constants /proc/self/cwd/external/iree_core/runtime/src/iree/hal/command_buffer.c:518:26
    #10 0x7f195e9d058c in iree_hal_module_command_buffer_push_constants /proc/self/cwd/external/iree_core/runtime/src/iree/modules/hal/module.c:779:10
    #11 0x7f195eb49c15 in iree_vm_shim_rriCiD_v /proc/self/cwd/external/iree_core/runtime/src/iree/vm/shims.c:56:1
    #12 0x7f195eb31d97 in iree_vm_native_module_issue_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/native_module.c:338:7
    #13 0x7f195eb309e9 in iree_vm_native_module_begin_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/native_module.c:392:10
    #14 0x7f195ea4632c in iree_vm_bytecode_issue_import_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:452:7
    #15 0x7f195ea41d2e in iree_vm_bytecode_call_import_variadic /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:609:10
    #16 0x7f195ea29d9d in iree_vm_bytecode_dispatch /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:1667:5
    #17 0x7f195e9ff564 in iree_vm_bytecode_dispatch_begin /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/dispatch.c:636:10
    #18 0x7f195e9f3fd0 in iree_vm_bytecode_module_begin_call /proc/self/cwd/external/iree_core/runtime/src/iree/vm/bytecode/module.c:779:10
    #19 0x7f195eb10c2e in iree_vm_begin_invoke /proc/self/cwd/external/iree_core/runtime/src/iree/vm/invocation.c:504:7
    #20 0x7f195eb0e64a in iree_vm_invoke /proc/self/cwd/external/iree_core/runtime/src/iree/vm/invocation.c:302:26
    #21 0x7f195e95abb8 in iree::pjrt::LoadedExecutableInstance::BatchExecute(PJRT_LoadedExecutable_Execute_Args*) /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1797:9
    #22 0x7f195e964d8d in iree::pjrt::LoadedExecutableInstance::BindApi(PJRT_Api*)::$_54::operator()(PJRT_LoadedExecutable_Execute_Args*) const /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1590:61
    #23 0x7f195e964d34 in iree::pjrt::LoadedExecutableInstance::BindApi(PJRT_Api*)::$_54::__invoke(PJRT_LoadedExecutable_Execute_Args*) /proc/self/cwd/iree/integrations/pjrt/common/api_impl.cc:1587:8
    #24 0x7f1b8e11a00a in xla::PjRtCApiLoadedExecutable::Execute(absl::lts_20230125::Span<std::vector<xla::PjRtBuffer*, std::allocator<xla::PjRtBuffer*> > const>, xla::ExecuteOptions const&, std::optional<std::vector<xla::PjRtFuture<absl::lts_20230125::Status>, std::allocator<xla::PjRtFuture<absl::lts_20230125::Status> > > >&) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0xcaf00a) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #25 0x7f1b906deca6 in xla::ifrt::PjRtLoadedExecutable::Execute(absl::lts_20230125::Span<tsl::RCReference<xla::ifrt::Array> >, xla::ExecuteOptions const&, std::optional<xla::ifrt::DeviceList>) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x3273ca6) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #26 0x7f1b8e0c8836 in absl::lts_20230125::StatusOr<xla::PyExecuteResults> xla::(anonymous namespace)::ExecuteShardedOnLocalDevicesInternal<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, xla::(anonymous namespace)::ShardedBufferAdapter<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >(xla::ExecuteOptions const&, std::shared_ptr<xla::PyClient> const&, xla::ifrt::LoadedExecutable*, absl::lts_20230125::Span<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > const>, std::optional<std::vector<xla::PjRtFuture<absl::lts_20230125::Status>, std::allocator<xla::PjRtFuture<absl::lts_20230125::Status> > > >&) py_executable.cc
    #27 0x7f1b8e0c9b8d in xla::PyLoadedExecutable::ExecuteSharded(std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0xc5eb8d) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #28 0x7f1b8dd9b1d3 in void pybind11::cpp_function::initialize<xla::ValueOrThrowWrapper<absl::lts_20230125::StatusOr<xla::PyExecuteResults> (std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), xla::PyLoadedExecutable>, xla::PyExecuteResults, xla::PyLoadedExecutable&, std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::arg, pybind11::arg_v>(xla::ValueOrThrowWrapper<absl::lts_20230125::StatusOr<xla::PyExecuteResults> (std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), xla::PyLoadedExecutable>&&, xla::PyExecuteResults (*)(xla::PyLoadedExecutable&, std::vector<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > >, std::allocator<std::variant<xla::PyArray, std::vector<xla::PyArray, std::allocator<xla::PyArray> > > > >, bool), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&)::'lambda1'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x9301d3) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)
    #29 0x7f1b8dd6fed0 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) (/usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so+0x904ed0) (BuildId: 8d2c019e2cf0e2ad84df858ea45f0237)

SUMMARY: AddressSanitizer: use-after-poison (/usr/lib/llvm-14/lib/clang/14.0.0/lib/linux/libclang_rt.asan-x86_64.so+0xccb99) (BuildId: 0fc20d2022c0d572c45850b6d559d9ccfabd5443) in __asan_memcpy
Shadow bytes around the buggy address:
  0x0c5a80087c30: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x0c5a80087c40: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x0c5a80087c50: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x0c5a80087c60: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x0c5a80087c70: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
=>0x0c5a80087c80:[f7]f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
  0x0c5a80087c90: f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
  0x0c5a80087ca0: f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
  0x0c5a80087cb0: f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
  0x0c5a80087cc0: f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
  0x0c5a80087cd0: f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7 f7
Shadow byte legend (one shadow byte represents 8 application bytes):
  Addressable:           00
  Partially addressable: 01 02 03 04 05 06 07 
  Heap left redzone:       fa
  Freed heap region:       fd
  Stack left redzone:      f1
  Stack mid redzone:       f2
  Stack right redzone:     f3
  Stack after return:      f5
  Stack use after scope:   f8
  Global redzone:          f9
  Global init order:       f6
  Poisoned by user:        f7
  Container overflow:      fc
  Array cookie:            ac
  Intra object redzone:    bb
  ASan internal:           fe
  Left alloca redzone:     ca
  Right alloca redzone:    cb
==12037==ABORTING