checkpoint/opt_state/array_metadatas/process_0

{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 64], "chunk_shape": [7, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 128], "chunk_shape": [1556, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 128], "chunk_shape": [20, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.temporal_enc", "write_shape": [6, 64], "chunk_shape": [6, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.type_enc", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 64], "chunk_shape": [7, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.kernel", "write_shape": [64, 7], "chunk_shape": [64, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_query", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.kernel", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.kernel", "write_shape": [71, 64], "chunk_shape": [71, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.kernel", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 64], "chunk_shape": [7, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 128], "chunk_shape": [1556, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 128], "chunk_shape": [20, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.temporal_enc", "write_shape": [6, 64], "chunk_shape": [6, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.type_enc", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 64], "chunk_shape": [7, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 8, 64], "chunk_shape": [8, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 8], "chunk_shape": [8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.kernel", "write_shape": [64, 8, 8], "chunk_shape": [64, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.kernel", "write_shape": [64, 7], "chunk_shape": [64, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_query", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.kernel", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.kernel", "write_shape": [71, 64], "chunk_shape": [71, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.kernel", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.bias", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.kernel", "write_shape": [64, 1, 64], "chunk_shape": [64, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.pos_emb", "write_shape": [5, 64], "chunk_shape": [5, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.1.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}