Publish PPO1 + AIM500 inference bundle

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/rollouts.mp4 filter=lfs diff=lfs merge=lfs -text
+models/ppo1-aim500/0/default/d/dd02eeee2256c0c61da40ade1ef8cf65 filter=lfs diff=lfs merge=lfs -text
+models/ppo1-aim500/0/default/ocdbt.process_0/d/0e705ab1649291dcd4833d7c54c84c68 filter=lfs diff=lfs merge=lfs -text
+models/ppo1-aim500/0/default/ocdbt.process_0/d/1842d714394728bf2cd8a101ca51be79 filter=lfs diff=lfs merge=lfs -text
+models/ppo1-aim500/0/default/ocdbt.process_0/d/76b87323ea6fece1f3000f3671f7a65c filter=lfs diff=lfs merge=lfs -text
+models/ppo1-aim500/0/default/ocdbt.process_0/d/e127cf6e78fa29c319f572ccfd0c1d80 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,151 @@
+---
+language:
+- en
+tags:
+- world-model
+- robotics
+- pusht
+- jax
+- latent-diffusion
+datasets:
+- smbml/future-sight-pusht
+inference: false
+---
+
+<p align="center"><img src="assets/wordmark.svg" alt="Future Sight" width="420"></p>
+
+# PushT world models
+
+Action-conditioned world models from S. Bateman et al., **Training Controllable
+World Models via Action Influence Maximization**, CoRL 2026.
+
+PushT provides a controlled wind tunnel for studying how data scale and
+composition affect world-model prediction and generalization. These models
+compare large passive collections with smaller, task-specific active exploration
+collections using a shared architecture. The paired
+[Future Sight PushT dataset](https://huggingface.co/datasets/smbml/future-sight-pusht)
+provides the training and evaluation data.
+
+<video controls autoplay loop muted playsinline width="768" src="https://huggingface.co/smbml/future-sight-pusht-models/resolve/main/assets/rollouts.mp4"></video>
+
+**Ground truth and AIM predictions**, with identical recorded actions. Top:
+behavior-cloning evaluation; bottom: PushAway evaluation. Each clip shows 33
+predicted transitions at 256 × 256 resolution. These examples illustrate model
+behavior; they are not an aggregate benchmark or a ranking of models.
+
+## Models at a glance
+
+**1 checkpoint · 256 × 256 RGB · approximately 580 MB each · 0.58 GB total**
+
+- **Inputs:** RGB observations and continuous 6D end-effector delta actions.
+- **Outputs:** predicted RGB frames at 256 × 256 resolution.
+- **Format:** native Orbax inference bundles, including the frozen VAE.
+- **Weights:** exponential moving average (EMA); no optimizer or training-resume state.
+
+## Choose a model
+
+`PPO1` and `Recovery1` use the first 1% of the corresponding training metadata;
+`PPO100` and `Recovery100` use the full corpus. The `500` suffix denotes the
+additional collection budget. AIM500 uses 100 trajectories from each of five
+collection iterations, with a 2:1 PPO-to-AIM sampling mixture.
+
+| Model | Training step | Download |
+| --- | ---: | ---: |
+| [PPO1 + AIM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500) | 36,000 | 580.1 MB |
+
+The **fixed reward estimate** ablation uses AIM data without reward
+re-estimation. Novelty-only and MI-only models ablate the reward components.
+Transferred AIM adds exploration collected from the PPO-based setting to a
+Recovery1 training mixture. See the experiment-specific recipes in Future Sight
+for exact subset and mixture definitions; the budget suffix alone does not
+specify sampling probabilities.
+
+## Download and load one model
+
+Run from a Future Sight checkout with its Pixi environment installed and one
+visible GPU, for example `CUDA_VISIBLE_DEVICES=0 pixi run python example.py`
+after saving the following snippet as `example.py`. The example
+uses existing `huggingface_hub` and Future Sight APIs. No Transformers model
+conversion or `trust_remote_code` is required.
+
+```python
+import os
+from pathlib import Path
+
+import jax
+from huggingface_hub import snapshot_download
+
+from future_sight.models.components.sharding import create_device_mesh
+from future_sight.models.loading import load_world_model
+
+cache = Path(os.environ.get("FUTURE_SIGHT_CACHE_DIR", "cache")).expanduser()
+snapshot = Path(snapshot_download(
+    repo_id="smbml/future-sight-pusht-models",
+    revision="main",  # Use a full HF commit hash to pin an experiment.
+    allow_patterns=["models.json", "models/ppo1-aim500/*"],
+    cache_dir=cache / "models",
+))
+
+jax.devices()  # Initialize JAX before any TensorFlow-backed data loader.
+mesh = create_device_mesh(data=1, model=1)
+model, manifest = load_world_model(str(snapshot / "models/ppo1-aim500"), mesh)
+print(manifest["weights"]["source_training_step"])
+```
+
+The example explicitly honors `FUTURE_SIGHT_CACHE_DIR`; the HF library itself
+does not interpret that Future Sight variable. Change the folder in
+`allow_patterns` and the load path to select another model. HF uses your existing
+login when available; these public files are also downloadable anonymously.
+Completed files are reused from the HF cache, and Orbax reads them in place.
+
+The restored model can be used with Future Sight's Python inference session,
+standalone PushT evaluator, or local interactive application. Model-specific HF
+CLI integration and a hosted interactive Space are not provided by this repo.
+
+## Architecture and provenance
+
+All models use the latent diffusion forcing B/2 architecture: 12 transformer
+layers, width 768, 12 attention heads, and 2 × 2 latent patches. The VAE maps
+256 × 256 images to 32 × 32 × 4 latents. The published manifests specify the
+exact action bounds, x-prediction / v-space training objective, BF16 runtime,
+128 inference steps, four-frame sequence window and three cached predecessors.
+
+Each bundle includes a versioned manifest with its training step, source commit,
+and hashes of every restored tensor. [models.json](models.json) additionally
+records downloadable file sizes and SHA-256 digests. Public metadata removes
+private storage paths and uses the `future_sight` import namespace; checkpoint
+payloads and tensor identities are unchanged. Historical source repository names
+remain in provenance. Private training-contract files are not distributed.
+
+The frozen StableVAE is derived from the `sd-vae-ft-mse` family. Upstream
+[Stability AI weights](https://huggingface.co/stabilityai/sd-vae-ft-mse) and the
+[Flax conversion](https://huggingface.co/pcuenq/sd-vae-ft-mse-flax) are marked MIT.
+The original upstream HF revision was not recorded; the bundle records the
+converted VAE's content identity. This upstream license statement does not assign
+a license to the separately trained Future Sight weights.
+
+## Validation and limitations
+
+All 16 source checkpoints passed strict restoration and finite-output checks on two selected evaluation cases each. They also ran in the local interactive model switcher. Published model files are listed above.
+
+The local checks used JAX/JAXLIB 0.10.2, Flax 0.12.7 and Orbax 0.12.1 on an
+NVIDIA RTX PRO 6000 Blackwell GPU. They do not establish compatibility with every
+hardware or dependency combination. Small local checks and videos do not reproduce
+the complete paper evaluation. Samples may drift or produce implausible dynamics,
+particularly over longer horizons or unfamiliar actions.
+
+## Citation
+
+If you use these models or the paired dataset, please cite:
+
+```bibtex
+@inproceedings{bateman2026aim,
+  title = {Training Controllable World Models via Action Influence Maximization},
+  author = {Bateman, Samuel M. and Yin, Tenny and Zheng, Chongyi and Huang, Lei
+            and Wang, Brian and Eysenbach, Ben and Fisac, Jaime Fern\'andez and Shah, Dhruv},
+  booktitle = {Conference on Robot Learning},
+  year = {2026}
+}
+```
+
+Paper link forthcoming.

assets/rollouts.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f334766610a6a51dbdc305dd3665610dadac832b803debde4e5cfa93d40b1cf
+size 127188

models.json

@@ -0,0 +1,88 @@
+{
+  "schema_version": 1,
+  "repo_id": "smbml/future-sight-pusht-models",
+  "dataset_repo_id": "smbml/future-sight-pusht",
+  "models": [
+    {
+      "id": "ppo1-aim500",
+      "label": "PPO1 + AIM500",
+      "path": "models/ppo1-aim500",
+      "state_root_sha256": "1edd44b711c344f9890be95d9daa530b9b72006b189a6aa2d741050d4e9dbef5",
+      "training_step": 36000,
+      "weights": "ema",
+      "size_bytes": 580115101,
+      "files": [
+        {
+          "path": "0/_CHECKPOINT_METADATA",
+          "size_bytes": 275,
+          "sha256": "b371aedea719086c98dbe4cb71e7eb25ba21d96b53b1cbe9ffa85455d33485d1"
+        },
+        {
+          "path": "0/default/_METADATA",
+          "size_bytes": 234269,
+          "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789"
+        },
+        {
+          "path": "0/default/_sharding",
+          "size_bytes": 166020,
+          "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d"
+        },
+        {
+          "path": "0/default/array_metadatas/process_0",
+          "size_bytes": 85571,
+          "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b"
+        },
+        {
+          "path": "0/default/d/dd02eeee2256c0c61da40ade1ef8cf65",
+          "size_bytes": 113007,
+          "sha256": "aaa1e489765032d2bdc075d8a37c83433370e35fa4142e6a37120a8053ec20cf"
+        },
+        {
+          "path": "0/default/manifest.ocdbt",
+          "size_bytes": 120,
+          "sha256": "b78789acce1eceeac3264df46c3da5ec785aa92f35b63e5f9f31ef1f05cbd6a5"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/0e705ab1649291dcd4833d7c54c84c68",
+          "size_bytes": 48422912,
+          "sha256": "8d245e8624e9123d4d637310346dbcfbcd2f9c5cc9d1905abfc248d704cc1668"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/1842d714394728bf2cd8a101ca51be79",
+          "size_bytes": 29741056,
+          "sha256": "7830cf618c4018c86ce52eb9ce01bdac7c97078c942ad0b8b9ea8c3b7048be7d"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/76b87323ea6fece1f3000f3671f7a65c",
+          "size_bytes": 491028480,
+          "sha256": "8bb949e4e7c33ec2f24c37b8a340e490b2d41d2b6d3739c1fd32d8cf677ecc0f"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/94c638ab267fd6b890565a0a207d7430",
+          "size_bytes": 298,
+          "sha256": "89e738033d1fe4e587398a0927a10a60cc62e47231564088ccd9a4423f0e11d2"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/d18ef7b0f7f577c5ed635340333bc744",
+          "size_bytes": 1490,
+          "sha256": "5b8fcd3783961e97e1633bdc51e49c949a1ce2a6b328dadf7bd1f95e13d42966"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/d/e127cf6e78fa29c319f572ccfd0c1d80",
+          "size_bytes": 10190848,
+          "sha256": "24591900946f31f3b8a7e2787ee81d4710600a58c420f8290567bfa4f805bbd4"
+        },
+        {
+          "path": "0/default/ocdbt.process_0/manifest.ocdbt",
+          "size_bytes": 370,
+          "sha256": "96b8e3d4974cf3623b3902fb2d7dd7a343ef60a7d448d190dbaedc756aaa5032"
+        },
+        {
+          "path": "checkpoint_manifest.json",
+          "size_bytes": 130385,
+          "sha256": "57078aa520fa28dba100d56593b1a42ab6c58b0c3b631d3bbdae929282d59cf5"
+        }
+      ]
+    }
+  ]
+}

models/ppo1-aim500/0/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225449825518755, "commit_timestamp_nsecs": 1786225450487340631, "custom_metadata": {}}

models/ppo1-aim500/0/default/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]}

models/ppo1-aim500/0/default/d/dd02eeee2256c0c61da40ade1ef8cf65

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aaa1e489765032d2bdc075d8a37c83433370e35fa4142e6a37120a8053ec20cf
+size 113007

models/ppo1-aim500/0/default/ocdbt.process_0/d/0e705ab1649291dcd4833d7c54c84c68

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d245e8624e9123d4d637310346dbcfbcd2f9c5cc9d1905abfc248d704cc1668
+size 48422912

models/ppo1-aim500/0/default/ocdbt.process_0/d/1842d714394728bf2cd8a101ca51be79

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7830cf618c4018c86ce52eb9ce01bdac7c97078c942ad0b8b9ea8c3b7048be7d
+size 29741056

models/ppo1-aim500/0/default/ocdbt.process_0/d/76b87323ea6fece1f3000f3671f7a65c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bb949e4e7c33ec2f24c37b8a340e490b2d41d2b6d3739c1fd32d8cf677ecc0f
+size 491028480

models/ppo1-aim500/0/default/ocdbt.process_0/d/e127cf6e78fa29c319f572ccfd0c1d80

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24591900946f31f3b8a7e2787ee81d4710600a58c420f8290567bfa4f805bbd4
+size 10190848