diff --git "a/jepa2_run_warmup/best/step_150000_20260311-122222-UTC/array_metadatas/process_0" "b/jepa2_run_warmup/best/step_150000_20260311-122222-UTC/array_metadatas/process_0"
new file mode 100644--- /dev/null
+++ "b/jepa2_run_warmup/best/step_150000_20260311-122222-UTC/array_metadatas/process_0"
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "ema_params.block_0.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.block_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.final_layer.Dense_0.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.final_layer.Dense_0.kernel", "write_shape": [768, 1536], "chunk_shape": [768, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.final_layer.Dense_1.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.final_layer.Dense_1.kernel", "write_shape": [768, 16], "chunk_shape": [768, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.patch_embed.Conv_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.patch_embed.Conv_0.kernel", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.t_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.t_embed.Dense_0.kernel", "write_shape": [256, 768], "chunk_shape": [256, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.t_embed.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.t_embed.Dense_1.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.y_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "ema_params.y_embed.Dense_0.kernel", "write_shape": [100, 768], "chunk_shape": [100, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.block_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.final_layer.Dense_0.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.final_layer.Dense_0.kernel", "write_shape": [768, 1536], "chunk_shape": [768, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.final_layer.Dense_1.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.final_layer.Dense_1.kernel", "write_shape": [768, 16], "chunk_shape": [768, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.patch_embed.Conv_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.patch_embed.Conv_0.kernel", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.t_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.t_embed.Dense_0.kernel", "write_shape": [256, 768], "chunk_shape": [256, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.t_embed.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.t_embed.Dense_1.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.y_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.y_embed.Dense_0.kernel", "write_shape": [100, 768], "chunk_shape": [100, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.block_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.final_layer.Dense_0.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.final_layer.Dense_0.kernel", "write_shape": [768, 1536], "chunk_shape": [768, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.final_layer.Dense_1.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.final_layer.Dense_1.kernel", "write_shape": [768, 16], "chunk_shape": [768, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.patch_embed.Conv_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.patch_embed.Conv_0.kernel", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.t_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.t_embed.Dense_0.kernel", "write_shape": [256, 768], "chunk_shape": [256, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.t_embed.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.t_embed.Dense_1.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.y_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.y_embed.Dense_0.kernel", "write_shape": [100, 768], "chunk_shape": [100, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.Dense_0.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.Dense_0.kernel", "write_shape": [768, 4608], "chunk_shape": [768, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MlpBlock_0.Dense_0.bias", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MlpBlock_0.Dense_0.kernel", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MlpBlock_0.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MlpBlock_0.Dense_1.kernel", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.block_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.final_layer.Dense_0.bias", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.final_layer.Dense_0.kernel", "write_shape": [768, 1536], "chunk_shape": [768, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.final_layer.Dense_1.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.final_layer.Dense_1.kernel", "write_shape": [768, 16], "chunk_shape": [768, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.patch_embed.Conv_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.patch_embed.Conv_0.kernel", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.t_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.t_embed.Dense_0.kernel", "write_shape": [256, 768], "chunk_shape": [256, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.t_embed.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.t_embed.Dense_1.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.y_embed.Dense_0.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.y_embed.Dense_0.kernel", "write_shape": [100, 768], "chunk_shape": [100, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "rng", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file