{"array_metadatas": [{"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.map.to_m_var", "write_shape": [4, 4], "chunk_shape": [4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_ranges", "write_shape": [22, 2], "chunk_shape": [22, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_transforms", "write_shape": [22, 6], "chunk_shape": [22, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.kinematics_transforms", "write_shape": [22, 7], "chunk_shape": [22, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.2.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}