{"array_metadatas": [{"array_metadata": {"param_name": "model.dynamics.action_up.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.action_up.kernel.value", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.diffusion_transformer.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.dynamics.timestep_embed.embedding.value", "write_shape": [25, 32], "chunk_shape": [25, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.lam.vq.drop.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "model.tokenizer.mask_patch.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.action_up.kernel.value", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.diffusion_transformer.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.dynamics.timestep_embed.embedding.value", "write_shape": [25, 32], "chunk_shape": [25, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.tokenizer.mask_patch.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.action_up.kernel.value", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.4.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.blocks.5.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.diffusion_transformer.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.dynamics.timestep_embed.embedding.value", "write_shape": [25, 32], "chunk_shape": [25, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_in.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.action_up.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.patch_up.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.lam.vq.codebook.value", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_dense.kernel.value", "write_shape": [32, 512], "chunk_shape": [32, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm1.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.decoder.output_dense.kernel.value", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.0.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.1.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.2.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense1.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_dense2.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.ffn_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.spatial_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.key.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.out.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.query.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.bias.value", "write_shape": [8, 64], "chunk_shape": [8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_attention.value.kernel.value", "write_shape": [512, 8, 64], "chunk_shape": [512, 8, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.blocks.3.temporal_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_dense.kernel.value", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.input_norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.bias.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.encoder.output_dense.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.tokenizer.mask_patch.value", "write_shape": [1, 1, 1, 768], "chunk_shape": [1, 1, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.2.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}